Cut Off Trees for Golf Event LeetCode solution in Java

1
class Solution {
2
  // approach: 1st store all the positions in the min heap acc. to their height
3
  // now start removing the elements from the heap and calculate their dis using bfs
4
  // if at any point we cann't reach at the next position return -1;
5
  // else keep on adding the distances and return;
6
  public int cutOffTree(List<List<Integer>> forest) {
7
    PriorityQueue<int[]> pq =
8
        new PriorityQueue<>((a, b) -> (forest.get(a[0]).get(a[1]) - forest.get(b[0]).get(b[1])));
9
    for (int i = 0; i < forest.size(); i++) {
10
      for (int j = 0; j < forest.get(0).size(); j++) {
11
        if (forest.get(i).get(j) > 1) pq.add(new int[] {i, j});
12
      }
13
    }
14
    int ans = 0;
15
    int curr[] = {0, 0};
16
    while (pq.size() > 0) {
17
      int[] temp = pq.poll();
18
      int dis = calcDis(forest, curr, temp);
19
      // System.out.println(dis+" "+temp.height);
20
      if (dis == -1) return -1;
21
      ans += dis;
22
      curr = temp;
23
    }
24
    return ans;
25
  }
26

27
  int calcDis(List<List<Integer>> forest, int start[], int end[]) {
28
    int n = forest.size(), m = forest.get(0).size();
29
    boolean vis[][] = new boolean[n][m];
30
    Queue<int[]> queue = new LinkedList<>();
31
    queue.add(start);
32
    vis[start[0]][start[1]] = true;
33
    int dis = 0;
34
    while (queue.size() > 0) {
35
      int len = queue.size();
36
      while (len-- > 0) {
37
        int temp[] = queue.remove();
38
        int r = temp[0], c = temp[1];
39
        if (r == end[0] && c == end[1]) return dis;
40
        if (r + 1 < n && !vis[r + 1][c] && forest.get(r + 1).get(c) != 0) {
41
          queue.add(new int[] {r + 1, c});
42
          vis[r + 1][c] = true;
43
        }
44
        if (r - 1 >= 0 && !vis[r - 1][c] && forest.get(r - 1).get(c) != 0) {
45
          queue.add(new int[] {r - 1, c});
46
          vis[r - 1][c] = true;
47
        }
48
        if (c - 1 >= 0 && !vis[r][c - 1] && forest.get(r).get(c - 1) != 0) {
49
          queue.add(new int[] {r, c - 1});
50
          vis[r][c - 1] = true;
51
        }
52
        if (c + 1 < m && !vis[r][c + 1] && forest.get(r).get(c + 1) != 0) {
53
          queue.add(new int[] {r, c + 1});
54
          vis[r][c + 1] = true;
55
        }
56
      }
57
      dis++;
58
    }
59
    return -1;
60
  }
61
}