Falling Squares LeetCode solution in Java

1
class Solution {
2
  public List<Integer> fallingSquares(int[][] positions) {
3
    SegmentNode root = new SegmentNode(0, Integer.MAX_VALUE, 0);
4
    List<Integer> ans = new ArrayList<>();
5
    int max = 0;
6
    for (int[] p : positions) {
7
      int left = p[0], height = p[1], right = left + height;
8
      int maxHeight = query(root, left, right) + height;
9
      max = Math.max(max, maxHeight);
10
      ans.add(max);
11
      add(root, left, right, maxHeight);
12
    }
13
    return ans;
14
  }
15

16
  public int query(SegmentNode root, int start, int end) {
17
    if (start <= root.start && end >= root.end) return root.maxHeight;
18
    if (start >= root.end || end <= root.start) return 0;
19
    if (root.left == null) return root.maxHeight;
20
    int mid = root.start + (root.end - root.start) / 2;
21
    if (end <= mid) {
22
      return query(root.left, start, end);
23
    } else if (start >= mid) {
24
      return query(root.right, start, end);
25
    }
26
    return Math.max(query(root.left, start, mid), query(root.right, mid, end));
27
  }
28

29
  public void add(SegmentNode root, int start, int end, int maxHeight) {
30
    if (start <= root.start && end >= root.end) {
31
      root.maxHeight = maxHeight;
32
      root.left = null;
33
      root.right = null;
34
      return;
35
    }
36
    if (start >= root.end || root.start >= end) return;
37
    if (root.left == null) {
38
      int mid = root.start + (root.end - root.start) / 2;
39
      root.left = new SegmentNode(root.start, mid, 0);
40
      root.right = new SegmentNode(mid, root.end, 0);
41
    }
42
    add(root.left, start, end, maxHeight);
43
    add(root.right, start, end, maxHeight);
44
    root.maxHeight = Math.max(root.left.maxHeight, root.right.maxHeight);
45
  }
46
}
47

48
class SegmentNode {
49
  public SegmentNode left, right;
50
  public int start, end, maxHeight;
51

52
  public SegmentNode(int start, int end, int maxHeight) {
53
    this.start = start;
54
    this.end = end;
55
    this.maxHeight = maxHeight;
56
  }
57
}