Kth smallest path xor sum

fn ch_alloc(ch: &mut [[u32; 2]], ct: &mut [u32], pool: &mut u32, free: &mut Vec<u32>) -> u32 {
  if let Some(id) = free.pop() {
    ch[id as usize] = [0, 0];
    ct[id as usize] = 0;
    id
  } else {
    let id = *pool;
    *pool += 1;
    id
  }
}

fn ch_merge(a: u32, b: u32, ch: &mut [[u32; 2]], ct: &mut [u32], free: &mut Vec<u32>) -> u32 {
  if b == 0 { return a; }
  if a == 0 { return b; }
  let [bl, br] = ch[b as usize];
  if bl == 0 && br == 0 { free.push(b); return a; }
  let [al, ar] = ch[a as usize];
  let nl = ch_merge(al, bl, ch, ct, free);
  let nr = ch_merge(ar, br, ch, ct, free);
  ch[a as usize] = [nl, nr];
  ct[a as usize] = if nl != 0 { ct[nl as usize] } else { 0 }
           + if nr != 0 { ct[nr as usize] } else { 0 };
  free.push(b);
  a
}

#[inline(always)]
fn ch_insert(root: u32, v: i32, ch: &mut [[u32; 2]], ct: &mut [u32], pool: &mut u32, free: &mut Vec<u32>) -> u32 {
  const BITS: usize = 17;
  let root = if root == 0 { ch_alloc(ch, ct, pool, free) } else { root };
  let mut path = [0u32; BITS + 1];
  path[0] = root;
  let mut cur = root;
  let mut is_new = false;
  for d in 0..BITS {
    let b = ((v >> (BITS - 1 - d)) & 1) as usize;
    if ch[cur as usize][b] == 0 {
      ch[cur as usize][b] = ch_alloc(ch, ct, pool, free);
      is_new = true;
    }
    cur = ch[cur as usize][b];
    path[d + 1] = cur;
  }
  if is_new { for &node in &path { ct[node as usize] += 1; } }
  root
}

fn ch_collect(root: u32, ch: &[[u32; 2]], out: &mut Vec<i32>, stk: &mut Vec<(u32, i32, u32)>) {
  if root == 0 { return; }
  stk.push((root, 0, 17));
  while let Some((s, sv, sb)) = stk.pop() {
    let (mut node, mut val, mut bits) = (s, sv, sb);
    loop {
      if bits == 0 { out.push(val); break; }
      bits -= 1;
      let bit = bits as i32;
      let [l, r] = ch[node as usize];
      if l != 0 && r != 0 { stk.push((r, val | (1 << bit), bits)); node = l; }
      else if l != 0       { node = l; }
      else                 { val |= 1 << bit; node = r; }
    }
  }
}

impl Solution {
  pub fn kth_smallest(par: Vec<i32>, vals: Vec<i32>, queries: Vec<Vec<i32>>) -> Vec<i32> {
    let n = par.len();
    let q = queries.len();

    let mut tree: Vec<Vec<usize>> = vec![vec![]; n];
    for i in 1..n { tree[par[i] as usize].push(i); }

    let mut query_at: Vec<Vec<(u32, usize)>> = vec![vec![]; n];
    for (i, qi) in queries.iter().enumerate() {
      query_at[qi[0] as usize].push((qi[1] as u32, i));
    }

    let mut ans = vec![0i32; q];
    let mut xp  = vec![0i32; n];
    xp[0] = vals[0];

    // ch[i]=[left,right], ct[i]=count; node 0 = null sentinel.
    // Peak live nodes <= 17*n (path nodes) + n (roots) = 18*n.
    let max_nodes = 18 * n + 1024;
    let mut ch: Vec<[u32; 2]> = vec![[0, 0]; max_nodes];
    let mut ct: Vec<u32>      = vec![0u32;   max_nodes];
    let mut pool: u32 = 1;
    let mut free: Vec<u32>    = Vec::with_capacity(4096);

    let mut roots: Vec<u32>              = vec![0u32; n];
    let mut col_stk: Vec<(u32, i32, u32)> = Vec::with_capacity(512);
    let mut col_buf: Vec<i32>             = Vec::with_capacity(512);

    let mut dfs: Vec<(usize, bool)> = vec![(0, false)];
    while let Some((node, done)) = dfs.pop() {
      if !done {
        dfs.push((node, true));
        for &c in &tree[node] {
          xp[c] = xp[node] ^ vals[c];
          dfs.push((c, false));
        }
      } else {
        // small-to-large: take child with largest ct as base
        let mut base = usize::MAX;
        let mut bcnt = 0u32;
        for &c in &tree[node] {
          let r = roots[c];
          if r != 0 && ct[r as usize] > bcnt { bcnt = ct[r as usize]; base = c; }
        }
        let mut root = if base != usize::MAX { roots[base] } else { 0 };
        root = ch_insert(root, xp[node], &mut ch, &mut ct, &mut pool, &mut free);
        for &c in &tree[node] {
          if c == base { continue; }
          let cr = roots[c];
          if cr != 0 { root = ch_merge(root, cr, &mut ch, &mut ct, &mut free); }
        }
        if !query_at[node].is_empty() {
          ch_collect(root, &ch, &mut col_buf, &mut col_stk);
          let sz = col_buf.len();
          for &(k, idx) in &query_at[node] {
            ans[idx] = if k as usize > sz { -1 } else { col_buf[k as usize - 1] };
          }
          col_buf.clear();
        }
        roots[node] = root;
      }
    }
    ans
  }
}