use std::collections::BinaryHeap;
use std::cmp::Reverse;

impl Solution {
    pub fn max_points(grid: Vec<Vec<i32>>, queries: Vec<i32>) -> Vec<i32> {
        let k = queries.len();

        // Store original indices of the queries
        let mut sorted_indices: Vec<usize> = (0..k).collect();
        sorted_indices.sort_unstable_by_key(|&i| queries[i]);

        let mut result = vec![0; k];
        let mut query_ptr = 0;
        // Callback closure to update result for all queries with value ≤ limit
        let mut set_answers = |limit: i32, count: i32| {
            while query_ptr < k && queries[sorted_indices[query_ptr]] <= limit {
                result[sorted_indices[query_ptr]] = count;
                query_ptr += 1;
            }
        };

        Self::bfs_with_threshold_callback(&grid, &mut set_answers);

        result
    }

    fn bfs_with_threshold_callback<F>(grid: &[Vec<i32>], callback: &mut F)
    where
        F: FnMut(i32, i32),
    {
        let directions = [(-1, 0), (1, 0), (0, -1), (0, 1)];
        let m = grid.len();
        let n = grid[0].len();

        let mut visited = vec![vec![false; n]; m];
        let mut heap = BinaryHeap::new();
        let mut count = 0;
        let mut limit = 0;

        heap.push(Reverse((grid[0][0], 0, 0)));
        visited[0][0] = true;

        while let Some(Reverse((value, row, col))) = heap.pop() {
            // Trigger callback when reaching a higher threshold
            if value > limit {
                callback(value, count);
                limit = value;
            }

            count += 1;

            for (dr, dc) in directions {
                let r = row as isize + dr;
                let c = col as isize + dc;
                if r >= 0 && r < m as isize && c >= 0 && c < n as isize {
                    let r = r as usize;
                    let c = c as usize;
                    if !visited[r][c] {
                        visited[r][c] = true;
                        heap.push(Reverse((grid[r][c], r, c)));
                    }
                }
            }
        }

        // Final callback to assign remaining unanswered queries
        callback(i32::MAX, count);
    }
}