Minimum cost path with teleportations

impl Solution {
  pub fn min_cost(grid: Vec<Vec<i32>>, k: i32) -> i32 {
    use std::collections::BinaryHeap;
    use std::cmp::Reverse;
    let m = grid.len();
    let n = grid[0].len();
    let k = k as usize;
    let total = m * n;

    // Sort flat cell indices by grid value (ascending).
    // sorted[rank] = flat index of the rank-th smallest cell.
    let mut sorted: Vec<usize> = (0..total).collect();
    sorted.sort_unstable_by_key(|&i| grid[i / n][i % n]);

    // enter_rank[flat] = highest rank whose value <= grid[flat/n][flat%n].
    // Initially each cell maps to its own rank, then we propagate backward
    // within tied-value groups so every cell gets the last rank of its group.
    let mut enter_rank = vec![0usize; total];
    for (rk, &idx) in sorted.iter().enumerate() {
      enter_rank[idx] = rk;
    }
    for i in (0..total - 1).rev() {
      let vi = grid[sorted[i]     / n][sorted[i]     % n];
      let vj = grid[sorted[i + 1] / n][sorted[i + 1] % n];
      if vi == vj {
        enter_rank[sorted[i]] = enter_rank[sorted[i + 1]];
      }
    }

    // State indices:
    //   Regular state (r, c, t): index = (r*n + c)*(k+1) + t
    //   Virtual hub  (rank, t):  index = total*(k+1) + rank*(k+1) + t
    let reg_states = total * (k + 1);
    let reg  = |r: usize, c: usize, t: usize| (r * n + c) * (k + 1) + t;
    let virt = |rk: usize, t: usize| reg_states + rk * (k + 1) + t;

    let inf = i32::MAX / 2;
    let mut dist = vec![inf; reg_states + reg_states /* virt reuses same size */];
    dist[reg(0, 0, 0)] = 0;

    let mut heap: BinaryHeap<Reverse<(i32, usize)>> = BinaryHeap::new();
    heap.push(Reverse((0, reg(0, 0, 0))));

    while let Some(Reverse((cost, state))) = heap.pop() {
      if cost > dist[state] { continue; }

      if state < reg_states {
        // ── Regular state ──
        let t  = state % (k + 1);
        let rc = state / (k + 1);
        let r  = rc / n;
        let c  = rc % n;

        // Normal move: right or down
        for (nr, nc) in [(r, c + 1), (r + 1, c)] {
          if nr < m && nc < n {
            let new_cost = cost + grid[nr][nc];
            let ns = reg(nr, nc, t);
            if new_cost < dist[ns] {
              dist[ns] = new_cost;
              heap.push(Reverse((new_cost, ns)));
            }
          }
        }

        // Enter teleport hub (consume one teleport).
        // Enter at the highest rank whose value <= grid[r][c].
        if t < k {
          let flat = r * n + c;
          let rk = enter_rank[flat];
          let ns = virt(rk, t + 1);
          if cost < dist[ns] {
            dist[ns] = cost;
            heap.push(Reverse((cost, ns)));
          }
        }
      } else {
        // ── Virtual hub state ──
        let rem = state - reg_states;
        let t   = rem % (k + 1);
        let rk  = rem / (k + 1);

        // Exit hub: teleport to the cell at this rank (cost 0).
        let flat = sorted[rk];
        let ns = reg(flat / n, flat % n, t);
        if cost < dist[ns] {
          dist[ns] = cost;
          heap.push(Reverse((cost, ns)));
        }

        // Descend hub chain: access cells with strictly smaller rank
        // (i.e., equal-or-smaller value).
        if rk > 0 {
          let ns = virt(rk - 1, t);
          if cost < dist[ns] {
            dist[ns] = cost;
            heap.push(Reverse((cost, ns)));
          }
        }
      }
    }

    (0..=k).map(|t| dist[reg(m - 1, n - 1, t)]).min().unwrap()
  }
}