1268. Search Suggestions System

• Time:
• Space:

C++JavaPython

struct TrieNode {
  vector<shared_ptr<TrieNode>> children;
  const string* word = nullptr;
  TrieNode() : children(26) {}
};

class Solution {
 public:
  vector<vector<string>> suggestedProducts(vector<string>& products,
                                           string searchWord) {
    vector<vector<string>> ans;

    for (const string& product : products)
      insert(product);

    shared_ptr<TrieNode> node = root;

    for (const char c : searchWord) {
      if (node == nullptr || node->children[c - 'a'] == nullptr) {
        node = nullptr;
        ans.push_back({});
        continue;
      }
      node = node->children[c - 'a'];
      ans.push_back(search(node));
    }

    return ans;
  }

 private:
  shared_ptr<TrieNode> root = make_shared<TrieNode>();

  void insert(const string& word) {
    shared_ptr<TrieNode> node = root;
    for (const char c : word) {
      const int i = c - 'a';
      if (node->children[i] == nullptr)
        node->children[i] = make_shared<TrieNode>();
      node = node->children[i];
    }
    node->word = &word;
  }

  vector<string> search(shared_ptr<TrieNode> node) {
    vector<string> res;
    dfs(node, res);
    return res;
  }

  void dfs(shared_ptr<TrieNode> node, vector<string>& ans) {
    if (ans.size() == 3)
      return;
    if (node == nullptr)
      return;
    if (node->word != nullptr)
      ans.push_back(*node->word);
    for (shared_ptr<TrieNode> child : node->children)
      dfs(child, ans);
  }
};

class TrieNode {
  public TrieNode[] children = new TrieNode[26];
  public String word;
}

class Solution {
  public List<List<String>> suggestedProducts(String[] products, String searchWord) {
    List<List<String>> ans = new ArrayList<>();

    for (final String product : products)
      insert(product);

    TrieNode node = root;

    for (final char c : searchWord.toCharArray()) {
      if (node == null || node.children[c - 'a'] == null) {
        node = null;
        ans.add(new ArrayList<>());
        continue;
      }
      node = node.children[c - 'a'];
      ans.add(search(node));
    }

    return ans;
  }

  private TrieNode root = new TrieNode();

  private void insert(final String word) {
    TrieNode node = root;
    for (final char c : word.toCharArray()) {
      final int i = c - 'a';
      if (node.children[i] == null)
        node.children[i] = new TrieNode();
      node = node.children[i];
    }
    node.word = word;
  }

  private List<String> search(TrieNode node) {
    List<String> res = new ArrayList<>();
    dfs(node, res);
    return res;
  }

  private void dfs(TrieNode node, List<String> ans) {
    if (ans.size() == 3)
      return;
    if (node == null)
      return;
    if (node.word != null)
      ans.add(node.word);
    for (TrieNode child : node.children)
      dfs(child, ans);
  }
}

class TrieNode:
  def __init__(self):
    self.children: dict[str, TrieNode] = {}
    self.word: str | None = None


class Solution:
  def suggestedProducts(self, products: list[str],
                        searchWord: str) -> list[list[str]]:
    ans = []
    root = TrieNode()

    def insert(word: str) -> None:
      node = root
      for c in word:
        node = node.children.setdefault(c, TrieNode())
      node.word = word

    def search(node: TrieNode | None) -> list[str]:
      res: list[str] = []
      dfs(node, res)
      return res

    def dfs(node: TrieNode | None, res: list[str]) -> None:
      if len(res) == 3:
        return
      if not node:
        return
      if node.word:
        res.append(node.word)
      for c in string.ascii_lowercase:
        if c in node.children:
          dfs(node.children[c], res)

    for product in products:
      insert(product)

    node = root

    for c in searchWord:
      if not node or c not in node.children:
        node = None
        ans.append([])
        continue
      node = node.children[c]
      ans.append(search(node))

    return ans