1632. Rank Transform of a Matrix

• Time: $O(mn\log mn)$
• Space: $O(mn)$

C++JavaPython

class UnionFind {
 public:
  void union_(int u, int v) {
    if (!id.contains(u))
      id[u] = u;
    if (!id.contains(v))
      id[v] = v;
    const int i = find(u);
    const int j = find(v);
    if (i != j)
      id[i] = j;
  }

  unordered_map<int, vector<int>> getGroupIdToValues() {
    unordered_map<int, vector<int>> groupIdToValues;
    for (const auto& [u, _] : id)
      groupIdToValues[find(u)].push_back(u);
    return groupIdToValues;
  }

 private:
  unordered_map<int, int> id;

  int find(int u) {
    return id[u] == u ? u : id[u] = find(id[u]);
  }
};

class Solution {
 public:
  vector<vector<int>> matrixRankTransform(vector<vector<int>>& matrix) {
    const int m = matrix.size();
    const int n = matrix[0].size();
    vector<vector<int>> ans(m, vector<int>(n));
    // {val: [(i, j)]}
    map<int, vector<pair<int, int>>> valToGrids;
    // rank[i] := the maximum rank of the row or column so far
    vector<int> maxRankSoFar(m + n);

    for (int i = 0; i < m; ++i)
      for (int j = 0; j < n; ++j)
        valToGrids[matrix[i][j]].emplace_back(i, j);

    for (const auto& [val, grids] : valToGrids) {
      UnionFind uf;
      for (const auto& [i, j] : grids)
        // Union i-th row with j-th col.
        uf.union_(i, j + m);
      for (const auto& [groupId, values] : uf.getGroupIdToValues()) {
        // Get the maximum rank of all the included rows and columns.
        int maxRank = 0;
        for (const int i : values)
          maxRank = max(maxRank, maxRankSoFar[i]);
        // Update all the rows and columns to maxRank + 1.
        for (const int i : values)
          maxRankSoFar[i] = maxRank + 1;
      }
      for (const auto& [i, j] : grids)
        ans[i][j] = maxRankSoFar[i];
    }

    return ans;
  }
};

class UnionFind {
  public void union(int u, int v) {
    id.putIfAbsent(u, u);
    id.putIfAbsent(v, v);
    final int i = find(u);
    final int j = find(v);
    if (i != j)
      id.put(i, j);
  }

  public Map<Integer, List<Integer>> getGroupIdToValues() {
    Map<Integer, List<Integer>> groupIdToValues = new HashMap<>();
    for (Map.Entry<Integer, Integer> entry : id.entrySet()) {
      final int u = entry.getKey();
      final int i = find(u);
      groupIdToValues.putIfAbsent(i, new ArrayList<>());
      groupIdToValues.get(i).add(u);
    }
    return groupIdToValues;
  }

  private Map<Integer, Integer> id = new HashMap<>();

  private int find(int u) {
    return id.getOrDefault(u, u) == u ? u : find(id.get(u));
  }
}

class Solution {
  public int[][] matrixRankTransform(int[][] matrix) {
    final int m = matrix.length;
    final int n = matrix[0].length;
    int[][] ans = new int[m][n];
    // {val: [(i, j)]}
    TreeMap<Integer, List<Pair<Integer, Integer>>> valToGrids = new TreeMap<>();
    // rank[i] := the maximum rank of the row or column so far
    int[] maxRankSoFar = new int[m + n];

    for (int i = 0; i < m; ++i)
      for (int j = 0; j < n; ++j) {
        final int val = matrix[i][j];
        valToGrids.putIfAbsent(val, new ArrayList<>());
        valToGrids.get(val).add(new Pair<>(i, j));
      }

    for (Map.Entry<Integer, List<Pair<Integer, Integer>>> entry : valToGrids.entrySet()) {
      final int val = entry.getKey();
      List<Pair<Integer, Integer>> grids = entry.getValue();
      UnionFind uf = new UnionFind();
      for (Pair<Integer, Integer> grid : grids) {
        final int i = grid.getKey();
        final int j = grid.getValue();
        // Union i-th row with j-th col.
        uf.union(i, j + m);
      }
      for (List<Integer> values : uf.getGroupIdToValues().values()) {
        // Get the maximum rank of all the included rows and columns.
        int maxRank = 0;
        for (final int i : values)
          maxRank = Math.max(maxRank, maxRankSoFar[i]);
        // Update all the rows and columns to maxRank + 1.
        for (final int i : values)
          maxRankSoFar[i] = maxRank + 1;
      }
      for (Pair<Integer, Integer> grid : grids) {
        final int i = grid.getKey();
        final int j = grid.getValue();
        ans[i][j] = maxRankSoFar[i];
      }
    }

    return ans;
  }
}

class UnionFind:
  def __init__(self):
    self.id = {}

  def union(self, u: int, v: int) -> None:
    self.id.setdefault(u, u)
    self.id.setdefault(v, v)
    i = self._find(u)
    j = self._find(v)
    if i != j:
      self.id[i] = j

  def getGroupIdToValues(self) -> dict[int, list[int]]:
    groupIdToValues = collections.defaultdict(list)
    for u in self.id.keys():
      groupIdToValues[self._find(u)].append(u)
    return groupIdToValues

  def _find(self, u: int) -> int:
    if self.id[u] != u:
      self.id[u] = self._find(self.id[u])
    return self.id[u]


class Solution:
  def matrixRankTransform(self, matrix: list[list[int]]) -> list[list[int]]:
    m = len(matrix)
    n = len(matrix[0])
    ans = [[0] * n for _ in range(m)]
    # {val: [(i, j)]}
    valToGrids = collections.defaultdict(list)
    # rank[i] := the maximum rank of the row or column so far
    maxRankSoFar = [0] * (m + n)

    for i, row in enumerate(matrix):
      for j, val in enumerate(row):
        valToGrids[val].append((i, j))

    for _, grids in sorted(valToGrids.items()):
      uf = UnionFind()
      for i, j in grids:
        # Union i-th row with j-th col.
        uf.union(i, j + m)
      for values in uf.getGroupIdToValues().values():
        # Get the maximum rank of all the included rows and columns.
        maxRank = max(maxRankSoFar[i] for i in values)
        for i in values:
          # Update all the rows and columns to maxRank + 1.
          maxRankSoFar[i] = maxRank + 1
      for i, j in grids:
        ans[i][j] = maxRankSoFar[i]

    return ans