3485. Longest Common Prefix of K Strings After Removal

• Time: $O(\Sigma |\texttt{words[i]}|)$
• Space: $O(\Sigma |\texttt{words[i]}|)$

C++JavaPython

struct TrieNode {
  vector<shared_ptr<TrieNode>> children;
  int count = 0;
  TrieNode() : children(26) {}
};

class Trie {
 public:
  Trie(int k) : k(k) {}

  void insert(const string& word) {
    shared_ptr<TrieNode> node = root;
    for (int i = 0; i < word.length(); ++i) {
      const int sz = i + 1;
      const int index = word[i] - 'a';
      if (node->children[index] == nullptr)
        node->children[index] = make_shared<TrieNode>();
      node = node->children[index];
      ++node->count;
      if (node->count >= k && prefixLengthsCount[sz]++ == 0)
        prefixLengths.insert(sz);
    }
  }

  void erase(const string& word) {
    shared_ptr<TrieNode> node = root;
    for (int i = 0; i < word.length(); ++i) {
      const int sz = i + 1;
      const int index = word[i] - 'a';
      if (node->children[index] == nullptr)
        node->children[index] = make_shared<TrieNode>();
      node = node->children[index];
      if (node->count == k && prefixLengthsCount[sz]-- == 1)
        prefixLengths.erase(sz);
      --node->count;
    }
  }

  int getLongestCommonPrefix() const {
    return prefixLengths.empty() ? 0 : *prefixLengths.begin();
  }

 private:
  const int k;
  shared_ptr<TrieNode> root = make_shared<TrieNode>();
  unordered_map<int, int> prefixLengthsCount;
  set<int, greater<>> prefixLengths;
};

class Solution {
 public:
  vector<int> longestCommonPrefix(vector<string>& words, int k) {
    vector<int> ans;
    Trie trie(k);

    for (const string& word : words)
      trie.insert(word);

    for (const string& word : words) {
      trie.erase(word);
      ans.push_back(trie.getLongestCommonPrefix());
      trie.insert(word);
    }

    return ans;
  }
};

class TrieNode {
  public TrieNode[] children = new TrieNode[26];
  public int count = 0;
}

class Trie {
  public Trie(int k) {
    this.k = k;
  }

  public void insert(final String word) {
    TrieNode node = root;
    for (int i = 0; i < word.length(); ++i) {
      final int sz = i + 1;
      final int index = word.charAt(i) - 'a';
      if (node.children[index] == null)
        node.children[index] = new TrieNode();
      node = node.children[index];
      ++node.count;
      if (node.count >= k && prefixLengthsCount.merge(sz, 1, Integer::sum) == 1)
        prefixLengths.add(sz);
    }
  }

  public void erase(final String word) {
    TrieNode node = root;
    for (int i = 0; i < word.length(); ++i) {
      final int sz = i + 1;
      final int index = word.charAt(i) - 'a';
      if (node.children[index] == null)
        node.children[index] = new TrieNode();
      node = node.children[index];
      if (node.count == k && prefixLengthsCount.merge(sz, -1, Integer::sum) == 0)
        prefixLengths.remove(sz);
      --node.count;
    }
  }

  public int getLongestCommonPrefix() {
    return prefixLengths.isEmpty() ? 0 : prefixLengths.first();
  }

  private final int k;
  private TrieNode root = new TrieNode();
  private Map<Integer, Integer> prefixLengthsCount = new HashMap<>();
  private TreeSet<Integer> prefixLengths = new TreeSet<>(Collections.reverseOrder());
}

class Solution {
  public int[] longestCommonPrefix(String[] words, int k) {
    final int[] ans = new int[words.length];
    Trie trie = new Trie(k);

    for (final String word : words)
      trie.insert(word);

    for (int i = 0; i < words.length; ++i) {
      trie.erase(words[i]);
      ans[i] = trie.getLongestCommonPrefix();
      trie.insert(words[i]);
    }

    return ans;
  }
}

class TrieNode:
  def __init__(self):
    self.children: dict[str, TrieNode] = {}
    self.count = 0


class Trie:
  def __init__(self, k: int):
    self.k = k
    self.root = TrieNode()
    self.prefixLengthsCount = collections.Counter()
    self.prefixLengths = SortedList()

  def insert(self, word: str) -> None:
    node = self.root
    for i, c in enumerate(word):
      sz = i + 1
      node = node.children.setdefault(c, TrieNode())
      node.count += 1
      if node.count >= self.k:
        self.prefixLengthsCount[sz] += 1
        if self.prefixLengthsCount[sz] == 1:
          self.prefixLengths.add(-sz)

  def erase(self, word: str) -> None:
    node = self.root
    for i, c in enumerate(word):
      sz = i + 1
      node = node.children[c]
      if node.count == self.k:
        self.prefixLengthsCount[sz] -= 1
        if self.prefixLengthsCount[sz] == 0:
          self.prefixLengths.remove(-sz)
      node.count -= 1

  def getLongestCommonPrefix(self) -> int:
    return 0 if not self.prefixLengths else -self.prefixLengths[0]


class Solution:
  def longestCommonPrefix(self, words: list[str], k: int) -> list[int]:
    ans = []
    trie = Trie(k)

    for word in words:
      trie.insert(word)

    for word in words:
      trie.erase(word)
      ans.append(trie.getLongestCommonPrefix())
      trie.insert(word)

    return ans