Most Stones Removed with Same Row or Column LeetCode solution in Java

1
class Solution {
2
  public int removeStones(int[][] stones) {
3
    int ret = 0;
4
    DisjointSet ds = new DisjointSet(stones.length);
5

6
    for (int i = 0; i < stones.length; i++) {
7
      for (int j = i + 1; j < stones.length; j++) {
8
        int s1[] = stones[i];
9
        int s2[] = stones[j];
10

11
        if (s1[0] == s2[0] || s1[1] == s2[1]) {
12
          ds.union(i, j);
13
        }
14
      }
15
    }
16
    // System.out.println(Arrays.toString(ds.sets));
17
    for (int i = 0; i < ds.sets.length; i++) {
18
      if (ds.sets[i] < 0) ret += Math.abs(ds.sets[i]) - 1;
19
    }
20

21
    return ret;
22
  }
23

24
  class DisjointSet {
25
    public int sets[];
26

27
    public DisjointSet(int size) {
28
      sets = new int[size];
29
      Arrays.fill(sets, -1);
30
    }
31

32
    //
33
    // weighted union
34
    // union->return size in negative
35
    public int union(int idx1, int idx2) {
36
      int p1 = find(idx1);
37
      int p2 = find(idx2);
38

39
      if (p1 == p2) { // same parent so directly returning size
40
        return sets[p1];
41
      } else {
42
        int w1 = Math.abs(sets[p1]);
43
        int w2 = Math.abs(sets[p2]);
44

45
        if (w1 > w2) {
46
          sets[p2] = p1;
47

48
          // collapsing FIND
49
          sets[idx1] = p1;
50
          sets[idx2] = p1;
51

52
          return sets[p1] = -(w1 + w2);
53
        } else {
54
          sets[p1] = p2;
55

56
          // collapsing FIND
57
          sets[idx1] = p2;
58
          sets[idx2] = p2;
59

60
          return sets[p2] = -(w1 + w2);
61
        }
62
      }
63
    }
64

65
    // collapsing FIND
66
    // find parent
67
    public int find(int idx) {
68
      int p = idx;
69
      while (sets[p] >= 0) {
70
        p = sets[p];
71
      }
72
      return p;
73
    }
74
  }
75
}