Find the Kth Smallest Sum of a Matrix With Sorted Rows LeetCode solution in Rust

1
use std::collections::{BinaryHeap, HashSet};
2

3
impl Solution {
4
  pub fn kth_smallest(mat: Vec<Vec<i32>>, k: i32) -> i32 {
5
    let mut heap = BinaryHeap::new();
6
    let mut set = HashSet::new();
7
    let width = mat[0].len() - 1;
8
    heap.push((
9
      -mat.iter().map(|row| row[0]).sum::<i32>(),
10
      vec![0; mat.len()],
11
    ));
12
    for _ in 1..k {
13
      let (cur, ind) = heap.pop().unwrap();
14
      for r in 0..ind.len() {
15
        if ind[r] < width {
16
          let mut tmp = ind.clone();
17
          tmp[r] += 1;
18
          if !set.contains(&tmp) {
19
            set.insert(tmp.clone());
20
            heap.push((cur + mat[r][tmp[r] - 1] - mat[r][tmp[r]], tmp));
21
          }
22
        }
23
      }
24
    }
25
    -heap.pop().unwrap().0
26
  }
27
}