给定一组字符串单词，找到字符串相等或一个字符串以另一个字符串开头的对的数量？

Question

我在做竞技编程时遇到了这个问题，我不太确定如何解决它最优。

给定一个字符串单词数组，找到字符串相等或一个字符串以另一个字符串开头的对的数量。换句话说，找到这样的对 i, j (0 ≤ i < j < words.length) that words[i] is a prefix of words[j], or words[j] is a prefix of words[i].

示例：

为了

words = ["back", "backdoor", "gammon", "backgammon", "comeback", "come", "door"]

输出应该是解决方案（单词）= 3。

相关对是：

words[0] = "back" and words[1] = "backdoor"
words[0] = "back" and words[3] = "backgammon"
words[4] = "comeback" and words[5] = "come"

下一个例子：

为了

words = ["abc", "a", "a", "b", "ab", "ac"]

输出应该是解决方案（单词）= 8。

相关对是：

words[0] = "abc" and words[1] = "a"
words[0] = "abc" and words[2] = "a"
words[0] = "abc" and words[4] = "ab"
words[1] = "a" and words[2] = "a"
words[1] = "a" and words[4] = "ab"
words[1] = "a" and words[5] = "ac"
words[2] = "a" and words[4] = "ab"
words[2] = "a" and words[5] = "ac"

我的解决方案（我认为这是幼稚的）在逻辑上工作，但会在更大的测试用例上进行 TLE。

def solution(words):
    res = 0
    n = len(words)
    for i in range(n):
        for j in range(i + 1, n):
            if words[i] == words[j] or words[i].startswith(words[j]) or words[j].startswith(words[i]):
                res += 1
        
    return res

我认为也许实现 Trie 或类似的方法可能很聪明......但我不知道如何实现它，坦率地说，我不确定它是否会将其优化到不会 TLE 的程度（？）

有什么想法吗？

Answer 1

您可以使用

itertools.combinations

创建所有可能对的生成器，然后对相关对进行求和

from itertools import combinations

print(sum(1 for w1, w2 in combinations(words, 2) if w1.startswith(w2) or w2.startswith(w1)))

如果你也想要这些话

pairs = [(w1, w2 )for w1, w2 in combinations(words, 2) if w1.startswith(w2) or w2.startswith(w1)]
print(pairs)
print(len(pairs))

输出

[('back', 'backdoor'), ('back', 'backgammon'), ('comeback', 'come')]
3