Skip to content

642. Design Search Autocomplete System 👍

  • Time:
    • Constructor: $O(\Sigma |\texttt{sentences[i]}|)$
    • input(c: chr): $O(1)$
  • Space: $O(\Sigma |\texttt{sentences[i]}|)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
class TrieNode implements Comparable<TrieNode> {
  public TrieNode[] children = new TrieNode[128];
  public String s = null;
  public int time = 0;
  public List<TrieNode> top3 = new ArrayList<>();

  public int compareTo(TrieNode o) {
    if (this.time == o.time)
      return this.s.compareTo(o.s);
    return o.time - this.time;
  }

  public void update(TrieNode node) {
    if (!this.top3.contains(node))
      this.top3.add(node);
    Collections.sort(top3);
    if (top3.size() > 3)
      top3.remove(top3.size() - 1);
  }
}

class AutocompleteSystem {
  public AutocompleteSystem(String[] sentences, int[] times) {
    for (int i = 0; i < sentences.length; ++i)
      insert(sentences[i], times[i]);
  }

  public List<String> input(char c) {
    if (c == '#') {
      insert(sb.toString(), 1);
      curr = root;
      sb = new StringBuilder();
      return new ArrayList<>();
    }

    sb.append(c);

    if (curr != null)
      curr = curr.children[c];
    if (curr == null)
      return new ArrayList<>();

    List<String> ans = new ArrayList<>();

    for (TrieNode node : curr.top3)
      ans.add(node.s);

    return ans;
  }

  private TrieNode root = new TrieNode();
  private TrieNode curr = root;
  private StringBuilder sb = new StringBuilder();

  private void insert(final String s, int time) {
    TrieNode node = root;
    for (final char c : s.toCharArray()) {
      if (node.children[c] == null)
        node.children[c] = new TrieNode();
      node = node.children[c];
    }
    node.s = s;
    node.time += time;

    // Walk the path again and update the node with leaf node
    TrieNode leaf = node;
    node = root;
    for (final char c : s.toCharArray()) {
      node = node.children[c];
      node.update(leaf);
    }
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
class TrieNode:
  def __init__(self):
    self.children: dict[str, TrieNode] = {}
    self.s: str | None = None
    self.time = 0
    self.top3: list[TrieNode] = []

  def __lt__(self, other):
    if self.time == other.time:
      return self.s < other.s
    return self.time > other.time

  def update(self, node) -> None:
    if node not in self.top3:
      self.top3.append(node)
    self.top3.sort()
    if len(self.top3) > 3:
      self.top3.pop()


class AutocompleteSystem:
  def __init__(self, sentences: list[str], times: list[int]):
    self.root = TrieNode()
    self.curr = self.root
    self.s: list[str] = []

    for sentence, time in zip(sentences, times):
      self._insert(sentence, time)

  def input(self, c: str) -> list[str]:
    if c == '#':
      self._insert(''.join(self.s), 1)
      self.curr = self.root
      self.s = []
      return []

    self.s.append(c)

    if self.curr:
      self.curr = self.curr.children.get(c, None)
    if not self.curr:
      return []
    return [node.s for node in self.curr.top3]

  def _insert(self, sentence: str, time: int) -> None:
    node = self.root
    for c in sentence:
      node = node.children.setdefault(c, TrieNode())
    node.s = sentence
    node.time += time

    leaf = node
    node: TrieNode = self.root
    for c in sentence:
      node = node.children[c]
      node.update(leaf)