Problem statement

https://leetcode.com/problems/prefix-and-suffix-search/

Solution 1

Let n be number of words, L be maximum length of word. One of the cool ways to solve this problem is the following: imagine, we have word apple, then what we need to find is substring suffix + # + prefix in string apple#apple. Now, quick way to find substring is for example create Trie with all suffixes of apple#apple, which include #.

Complexity

It is O(nL^2) time and space for init: for each word we need O(L^2) processing time. For f we need O(L) time and O(1) space, because we just traverse through tree. So, final complexity is O(nL^2 + QL), where Q is number of queries.

Code

class TrieNode:
    def __init__(self):
        self.children = {}
        self.index = 0

class Trie:
    def __init__(self):
        self.root = TrieNode()

    def insert(self, word, index): #None
        root = self.root
        root.index = index
        for symbol in word:
            root = root.children.setdefault(symbol, TrieNode())
            root.index = index
        
    def startsWith(self, word):
        root = self.root
        for symbol in word:
            if symbol not in root.children:
                return -1
            root = root.children[symbol]
        return root.index  


class WordFilter:
    def __init__(self, words):
        self.trie = Trie()
        self.words = {}
        
        for index, word in enumerate(words):
            long = word + "#" + word
            for i in range(len(word)):
                self.trie.insert(long[i:], index)
                
    def f(self, prefix, suffix):
        return self.trie.startsWith(suffix + "#" + prefix)

Solution 2

Finally, there is again solution with O(nL^3) time complexity, but without Tries at all: we just generate for each word all O(L^2) possible prefix/suffix pairs and put them into dictionary. It is working the best in practice, even though complexity is higher. This is because L is quite small and trie is quite heavy data structure with all connections, whereas here we keep only strings.

Complexity

Time is O(nL^3 + QL), space is O(nL^3).

Code

class WordFilter:
    def __init__(self, words):
        self.d = {}
        for i, word in enumerate(words):
            for p, s in product(range(len(word) + 1), repeat=2):
                self.d[word[:p], word[s:]] = i

    def f(self, prefix, suffix):
        return self.d.get((prefix, suffix), -1)