Longest common prefix of k strings after removal

impl Solution {
  pub fn longest_common_prefix(words: Vec<String>, k: i32) -> Vec<i32> {
    let n = words.len();
    let k = k as usize;
    let mut ans = vec![0i32; n];

    if n <= k { return ans; }

    let total_len: usize = words.iter().map(|w| w.len()).sum();
    let max_depth = words.iter().map(|w| w.len()).max().unwrap_or(0);

    // Build a Trie. Each node stores a count of words passing through it.
    // node 0 = root (virtual, not counted).
    let cap = total_len + 1;
    let mut children = vec![[0usize; 26]; cap];
    let mut count = vec![0usize; cap];
    let mut next_id = 1usize;

    // word_nodes[i][d] = trie node ID at depth d+1 for words[i].
    // Avoids re-traversing the trie and eliminates O(L) hash cost per level.
    let mut word_nodes: Vec<Vec<usize>> = Vec::with_capacity(n);

    for w in &words {
      let mut cur = 0;
      let mut path = Vec::with_capacity(w.len());
      for b in w.bytes() {
        let c = (b - b'a') as usize;
        if children[cur][c] == 0 {
          children[cur][c] = next_id;
          next_id += 1;
        }
        cur = children[cur][c];
        count[cur] += 1;
        path.push(cur);
      }
      word_nodes.push(path);
    }

    // For each depth, track the top-2 prefix counts and the node ID of the leader.
    // Used to answer: "is there another prefix of this length with count >= k?"
    let mut top1_cnt = vec![0usize; max_depth + 1];
    let mut top1_id  = vec![0usize; max_depth + 1];
    let mut top2_cnt = vec![0usize; max_depth + 1];

    for path in &word_nodes {
      for (d, &node) in path.iter().enumerate() {
        let l = d + 1;
        let c = count[node];
        if c > top1_cnt[l] {
          top2_cnt[l] = top1_cnt[l];
          top1_cnt[l] = c;
          top1_id[l]  = node;
        } else if node != top1_id[l] && c > top2_cnt[l] {
          top2_cnt[l] = c;
        }
      }
    }

    // suffix_max_l[l] = max depth l' in [l, max_depth] where top1_cnt[l'] >= k.
    // For "Part 1": prefixes longer than wi cannot involve word i, so global counts
    // are unaffected by its removal — suffix_max_l is valid as-is.
    let mut suffix_max_l = vec![0usize; max_depth + 2];
    let mut cur_max = 0usize;
    for l in (1..=max_depth).rev() {
      if top1_cnt[l] >= k { cur_max = cur_max.max(l); }
      suffix_max_l[l] = cur_max;
    }

    for i in 0..n {
      let wi_len = word_nodes[i].len();

      // Part 1: best length from prefixes strictly longer than wi.
      // word i has no such prefix, so removing it changes nothing at those depths.
      let best_longer = if wi_len + 1 <= max_depth { suffix_max_l[wi_len + 1] } else { 0 };

      // Part 2: scan depths 1..=wi_len (backwards for early exit).
      // At each depth we remove word i's contribution and check feasibility.
      let mut best_shorter = 0usize;
      for (d, &node) in word_nodes[i].iter().enumerate().rev() {
        let l = d + 1;
        let my_cnt = count[node];

        // Condition A: this prefix still has >= k words after removing word i.
        if my_cnt >= k + 1 {
          best_shorter = l;
          break;
        }

        // Condition B: some OTHER trie node at this depth has count >= k.
        let other_best = if top1_id[l] == node { top2_cnt[l] } else { top1_cnt[l] };
        if other_best >= k {
          best_shorter = l;
          break;
        }
      }

      ans[i] = best_longer.max(best_shorter) as i32;
    }

    ans
  }
}