Rewrite: Trie with arena allocation

use std::collections::HashMap;

#[derive(Default)]
struct TrieNode { children: HashMap<char, Box<TrieNode>>, end: bool }

pub struct Trie { root: TrieNode }

impl Trie {
    pub fn new() -> Self { Self { root: TrieNode::default() } }
    pub fn insert(&mut self, word: &str) {
        let mut cur = &mut self.root;
        for c in word.chars() {
            cur = cur.children.entry(c).or_default();
        }
        cur.end = true;
    }
    pub fn contains(&self, word: &str) -> bool {
        self.walk(word).map_or(false, |n| n.end)
    }
    pub fn starts_with(&self, prefix: &str) -> bool {
        self.walk(prefix).is_some()
    }
    fn walk(&self, s: &str) -> Option<&TrieNode> {
        let mut cur = &self.root;
        for c in s.chars() {
            cur = cur.children.get(&c)?;
        }
        Some(cur)
    }
}

class Trie {
  private root: { children: Map<string, any>; end: boolean } = { children: new Map(), end: false };
  insert(word: string): void {
    let cur = this.root;
    for (const c of word) {
      if (!cur.children.has(c)) cur.children.set(c, { children: new Map(), end: false });
      cur = cur.children.get(c);
    }
    cur.end = true;
  }
  contains(word: string): boolean {
    const n = this.walk(word);
    return n !== null && n.end;
  }
  startsWith(prefix: string): boolean { return this.walk(prefix) !== null; }
  private walk(s: string): { children: Map<string, any>; end: boolean } | null {
    let cur = this.root;
    for (const c of s) {
      const nx = cur.children.get(c);
      if (!nx) return null;
      cur = nx;
    }
    return cur;
  }
}

#[derive(Clone)]
struct TrieNode {
    children: [i32; 26],
    end: bool,
}

impl TrieNode {
    fn new() -> Self { Self { children: [-1; 26], end: false } }
}

pub struct Trie {
    nodes: Vec<TrieNode>,
}

impl Trie {
    pub fn new() -> Self {
        Self { nodes: vec![TrieNode::new()] }
    }

    pub fn insert(&mut self, word: &str) {
        let mut cur: usize = 0;
        for b in word.bytes() {
            let idx = (b - b'a') as usize;
            let next = self.nodes[cur].children[idx];
            cur = if next == -1 {
                self.nodes.push(TrieNode::new());
                let new_id = (self.nodes.len() - 1) as i32;
                self.nodes[cur].children[idx] = new_id;
                new_id as usize
            } else {
                next as usize
            };
        }
        self.nodes[cur].end = true;
    }

    pub fn contains(&self, word: &str) -> bool {
        self.walk(word).map_or(false, |id| self.nodes[id].end)
    }

    pub fn starts_with(&self, prefix: &str) -> bool {
        self.walk(prefix).is_some()
    }

    fn walk(&self, s: &str) -> Option<usize> {
        let mut cur: usize = 0;
        for b in s.bytes() {
            let idx = (b - b'a') as usize;
            let nx = self.nodes[cur].children[idx];
            if nx == -1 { return None; }
            cur = nx as usize;
        }
        Some(cur)
    }
}

// Arena-backed trie. Each node is 26 child indices + an end flag.
// Indices are stored in a single flat Int32Array for cache locality.
class Trie {
  private children: number[][] = [new Array(26).fill(-1)];
  private end: boolean[] = [false];

  insert(word: string): void {
    let cur = 0;
    for (let i = 0; i < word.length; i++) {
      const idx = word.charCodeAt(i) - 97;
      if (this.children[cur][idx] === -1) {
        this.children.push(new Array(26).fill(-1));
        this.end.push(false);
        this.children[cur][idx] = this.children.length - 1;
      }
      cur = this.children[cur][idx];
    }
    this.end[cur] = true;
  }

  contains(word: string): boolean {
    const id = this.walk(word);
    return id !== -1 && this.end[id];
  }

  startsWith(prefix: string): boolean {
    return this.walk(prefix) !== -1;
  }

  private walk(s: string): number {
    let cur = 0;
    for (let i = 0; i < s.length; i++) {
      const idx = s.charCodeAt(i) - 97;
      const nx = this.children[cur][idx];
      if (nx === -1) return -1;
      cur = nx;
    }
    return cur;
  }
}