1
class Solution {
2
  int[] parent;
3

4
  public int[][] matrixRankTransform(int[][] matrix) {
5
    int m = matrix.length;
6
    int n = matrix[0].length;
7
    int[][] answer = new int[m][n];
8

9
    // GROUP BY MATRIX VAL -> {X,Y}
10
    TreeMap<Integer, List<int[]>> map = new TreeMap<>();
11
    for (int i = 0; i < m; i++) {
12
      for (int j = 0; j < n; j++) {
13
        int[] xy = {i, j};
14
        int val = matrix[i][j];
15
        if (map.get(val) == null) map.put(val, new ArrayList<>());
16
        map.get(val).add(xy);
17
      }
18
    }
19

20
    // INITIALIZE MIN-RANK ARRAY FOR EVERY COL/ROW
21
    int[] minX = new int[m];
22
    int[] minY = new int[n];
23

24
    for (Integer key : map.keySet()) {
25
      List<int[]> list = map.get(key);
26

27
      // SPLIT TO GROUPS USING UNION FIND FOR VALs IN SAME COL/ROW
28
      int lSize = list.size();
29
      parent = new int[lSize];
30
      for (int i = 0; i < lSize; i++) parent[i] = i;
31

32
      // Group the xy by col and row then union by row & by col
33
      HashMap<Integer, List<Integer>> xMap = new HashMap<>();
34
      HashMap<Integer, List<Integer>> yMap = new HashMap<>();
35
      for (int i = 0; i < lSize; i++) {
36
        int[] xy = list.get(i);
37
        int x = xy[0];
38
        int y = xy[1];
39

40
        if (xMap.get(x) == null) xMap.put(x, new ArrayList<>());
41
        if (yMap.get(y) == null) yMap.put(y, new ArrayList<>());
42
        xMap.get(x).add(i);
43
        yMap.get(y).add(i);
44
      }
45

46
      // union by X
47
      for (Integer xKey : xMap.keySet()) {
48
        List<Integer> xList = xMap.get(xKey);
49
        for (int i = 1; i < xList.size(); i++) {
50
          union(xList.get(i - 1), xList.get(i));
51
        }
52
      }
53

54
      // union by Y
55
      for (Integer yKey : yMap.keySet()) {
56
        List<Integer> yList = yMap.get(yKey);
57
        for (int i = 1; i < yList.size(); i++) {
58
          union(yList.get(i - 1), yList.get(i));
59
        }
60
      }
61

62
      HashMap<Integer, List<int[]>> group = new HashMap<>();
63
      for (int i = 0; i < lSize; i++) {
64
        int grp = find(i);
65
        if (group.get(grp) == null) group.put(grp, new ArrayList<>());
66
        group.get(grp).add(list.get(i));
67
      }
68

69
      // SET ANSWER FOR EACH GROUP
70
      for (Integer grpKey : group.keySet()) {
71
        int max = 1;
72
        List<int[]> sublist = group.get(grpKey);
73

74
        // FIND MAX-RANK FOR THIS GROUP
75
        for (int[] xy : sublist) {
76
          int x = xy[0];
77
          int y = xy[1];
78

79
          max = Math.max(max, Math.max(minX[x], minY[y]));
80
        }
81

82
        // UPDATE ANSWER = MAX-RANK AND SET NEW MIN-RANK FOR ROW/COL = MAX-RANK+1
83
        for (int[] xy : sublist) {
84
          int x = xy[0];
85
          int y = xy[1];
86
          answer[x][y] = max;
87
          minX[x] = max + 1;
88
          minY[y] = max + 1;
89
        }
90
      }
91
    }
92
    return answer;
93
  }
94

95
  // UNION FIND IMPL
96
  void union(int a, int b) {
97
    int pa = find(a);
98
    int pb = find(b);
99
    parent[pb] = pa;
100
  }
101

102
  int find(int a) {
103
    int pa = parent[a];
104
    if (pa != a) {
105
      parent[a] = find(pa);
106
      return parent[a];
107
    } else return a;
108
  }
109
}