LeetCode 1268: Search Suggestions System

link

The question is a bit misleading. From the description, it feels like the lexicographical sorting should be applicable only if there are more than three products, which is not the case. As long as there are more than one suggestions, we must emit them in lexicographic order capping at three suggestions.

Sort and search

We can sort the products and for each prefix take the first three products that starts with that prefix.

Say |\texttt{products}| = n, longest product name has length m, and |\texttt{searchWord}| = k. Sorting has n \cdot \lg{n} comparisons and each comparison takes time m. So, sorting’s time: \mathcal{O}(m \cdot n \cdot \lg{n}). There are k searches through the products and each startswith takes time k.

Total time: \mathcal{O}(m \cdot n \cdot \lg{n} + n \cdot k^2). Space: \mathcal{O}(m \cdot \lg{n} + k).

class Solution:
    def suggestedProducts(
        self, products: List[str], searchWord: str
    ) -> List[List[str]]:
        products.sort()

        suggested_products = []
        for i in range(len(searchWord)):
            prefix = searchWord[: i + 1]
            suggestions = []
            for product in products:
                if len(suggestions) == 3:
                    break
                if product.startswith(prefix):
                    suggestions.append(product)

            suggested_products.append(suggestions)

        return suggested_products

Lexicographically sorted, the products with a common prefix should be consecutive. So, we can use binary search to find the insert position of a prefix. The three products starting from the insert index are the candidates.

Sorting’s time: \mathcal{O}(m \cdot n \cdot \lg{n}). Finding a prefix in the sorted \texttt{products} takes time k \cdot \lg{n}.

Total time: \mathcal{O}(m \cdot n \cdot \lg{n} + \lg{n} \cdot k^2). Space: \mathcal{O}(m \cdot \lg{n} + k).

class Solution:
    def find_insert_index(self, prefix, products) -> int:
        lo, hi = 0, len(products) - 1
        while lo <= hi:
            mid = (lo + hi) // 2
            if products[mid] < prefix:
                lo = mid + 1
            else:
                hi = mid - 1

        return lo

    def suggestedProducts(
        self, products: List[str], searchWord: str
    ) -> List[List[str]]:
        products.sort()

        suggested_products = []
        for i in range(len(searchWord)):
            prefix = searchWord[: i + 1]
            prefix_pos = self.find_insert_index(prefix, products)
            suggested_products.append(
                [
                    p
                    for p in products[prefix_pos : prefix_pos + 3]
                    if p.startswith(prefix)
                ]
            )

        return suggested_products

Since prefixes are from the same searchWord extending by one character in each iteration, we can use the insert index of the previous prefix as the lo for the binary search for the current prefix.

Total time: \mathcal{O}(m \cdot n \cdot \lg{n} + \lg{n} \cdot k^2). Space: \mathcal{O}(m \cdot \lg{n} + k).

class Solution:
    def find_insert_index(self, prefix, products, lo) -> int:
        hi = len(products) - 1
        while lo <= hi:
            mid = (lo + hi) // 2
            if products[mid] < prefix:
                lo = mid + 1
            else:
                hi = mid - 1

        return lo

    def suggestedProducts(
        self, products: List[str], searchWord: str
    ) -> List[List[str]]:
        products.sort()

        suggested_products = []
        prefix_pos = 0
        for i in range(len(searchWord)):
            prefix = searchWord[: i + 1]
            prefix_pos = self.find_insert_index(prefix, products, prefix_pos)
            suggested_products.append(
                [
                    p
                    for p in products[prefix_pos : prefix_pos + 3]
                    if p.startswith(prefix)
                ]
            )

        return suggested_products

Trie

We can use a trie to find product names with a given prefix.

class TrieNode:
    def __init__(self):
        self.end_of_word = False
        self.edges = {}

class Trie:
    def __init__(self):
        self.root = TrieNode()
    
    def insert(self, word):
        u = self.root
        for char in word:
            if char not in u.edges:
                u.edges[char] = TrieNode()
            u = u.edges[char]
        u.end_of_word = True

    def __get(self, word):
        u = self.root
        for char in word:
            if char not in u.edges:
                return None
            u = u.edges[char]
        return u

    def __collect(self, u, curr_word, words):
        if u.end_of_word:
            words.add("".join(curr_word))
        for v_char, v in u.edges.items():
            curr_word.append(v_char)
            self.__collect(v, curr_word, words)
            curr_word.pop()
    
    def words_with_prefix(self, word):
        node = self.__get(word)
        if node is None:
            return set()
        suffixes = set()
        self.__collect(node, [], suffixes)
        return [word+suff for suff in suffixes]

class Solution:
    def suggestedProducts(self, products: List[str], searchWord: str) -> List[List[str]]:
        trie = Trie()
        for p in products:
            trie.insert(p)

        suggested_products = []
        for i in range(len(searchWord)):
            prefix = searchWord[:i+1]
            suggestions = trie.words_with_prefix(prefix)
            if not suggestions:
                suggested_products.append([])
                continue

            suggestions.sort()
            if len(suggestions) <= 3:
                suggested_products.append(suggestions)
            else:
                suggested_products.append(suggestions[:3])
        
        return suggested_products

We could visit the edges in lexicographically sorted order and limit the number of words we collect.

Inserting n products into the trie takes time m \cdot n. Each of the k prefix searches takes time m.

Total time: \mathcal{O}(m \cdot n + k \cdot m). Space: \mathcal{O}(m \cdot n + k).

class TrieNode:
    def __init__(self):
        self.end_of_word = False
        self.edges = {}

class Trie:
    def __init__(self):
        self.root = TrieNode()
    
    def insert(self, word):
        u = self.root
        for char in word:
            if char not in u.edges:
                u.edges[char] = TrieNode()
            u = u.edges[char]
        u.end_of_word = True

    def __get(self, word):
        u = self.root
        for char in word:
            if char not in u.edges:
                return None
            u = u.edges[char]
        return u

    def __collect(self, u, curr_word, words, limit):
        if len(words) == limit:
            return

        if u.end_of_word:
            words.append("".join(curr_word))

        for v_char in sorted(u.edges):
            curr_word.append(v_char)
            self.__collect(u.edges[v_char], curr_word, words, limit)
            curr_word.pop()
    
    def words_with_prefix(self, word, limit):
        node = self.__get(word)
        if node is None:
            return []
        suffixes = []
        self.__collect(node, [], suffixes, limit)
        return [word+suff for suff in suffixes]

class Solution:
    def suggestedProducts(self, products: List[str], searchWord: str) -> List[List[str]]:
        trie = Trie()
        for p in products:
            trie.insert(p)

        suggested_products = []
        for i in range(len(searchWord)):
            prefix = searchWord[:i+1]
            suggestions = trie.words_with_prefix(prefix, 3)
            suggested_products.append(suggestions)
        
        return suggested_products

We can also use combination of trie and binary search. All suggestions must be subsets of the set of products with the prefix searchWord[0]. So, from the trie, we find this domain of suggestions — lexicographically sorted. Then, using binary search in this domain, we can find the suggestions for each prefix.

class TrieNode:
    def __init__(self):
        self.end_of_word = False
        self.edges = {}

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

    def insert(self, word):
        u = self.root
        for char in word:
            if char not in u.edges:
                u.edges[char] = TrieNode()
            u = u.edges[char]
        u.end_of_word = True

    def __get(self, word):
        u = self.root
        for char in word:
            if char not in u.edges:
                return None
            u = u.edges[char]
        return u

    def __collect(self, u, curr_word, words, limit):
        if len(words) == limit:
            return

        if u.end_of_word:
            words.append("".join(curr_word))

        for v_char in sorted(u.edges):
            curr_word.append(v_char)
            self.__collect(u.edges[v_char], curr_word, words, limit)
            curr_word.pop()

    def words_with_prefix(self, word, limit):
        node = self.__get(word)
        if node is None:
            return []
        suffixes = []
        self.__collect(node, [], suffixes, limit)
        return [word + suff for suff in suffixes]


class Solution:
    def find_insert_position(self, prefix, domain) -> int:
        lo, hi = 0, len(domain) - 1
        while lo <= hi:
            mid = (lo + hi) // 2
            if domain[mid] < prefix:
                lo = mid + 1
            else:
                hi = mid - 1

        return lo

    def suggestedProducts(
        self, products: List[str], searchWord: str
    ) -> List[List[str]]:
        trie = Trie()
        for p in products:
            trie.insert(p)

        suggestion_domain = trie.words_with_prefix(searchWord[0], len(products))
        suggested_products = []
        for pref_len in range(1, len(searchWord)+1):
            prefix = searchWord[: pref_len]
            prefix_pos = self.find_insert_position(prefix, suggestion_domain)
            suggested_products.append(
                [
                    p
                    for p in suggestion_domain[prefix_pos:prefix_pos+3]
                    if p.startswith(prefix)
                ]
            )

        return suggested_products

Leave a comment