Problem statement

https://leetcode.com/problems/rank-transform-of-a-matrix/

Solution

Quite difficult problem. Let us first start with the smallest number in our matrix. In general we create defaultict d, which for every value contain coordinates of cells with this value. We can notice, that we can start to work with small groups, then when we have next group we need to only care about what ranks we already put in out matrix. Imagine, than we first filled places where number x1 was, how we can do it? Actually at the first step we can put them all equal to 1. However on the next step when we consider number x2 > x1 we already have some ranks in our table, so we need to be carefull. For all numbers x2 in our table we can create graph, where nodes are n + m columns and rows and we have edge between row i and column j if we have value x2 at the intersection of these row and column. Then our graph can be separated into several connected components: where for each connected component we must have equal values. Different connected components can have different value, because we already put some ranks when we worked with x1.

  1. rank is variable for biggest rank so far in each n rows and m columns.
  2. Let us use Union Find data structure for efficient working with connected components. There is n rows and m columns, we will deal with point (edge in graph) (i, j) as i, j + n connection. We will use sparse dsu structure, that is we can have e.g. elements (1, 3, 7, 12, 20) in it, not neccesarily range from 1 to 20.
  3. For each edge, create connection: for i, j in d[a]: dsu.union(i, j + n)
  4. Now, we need to find connected components, use groups function inside our dsu structure. For each connected component we need to look at ranks we have: we need to take the biggest among them and add 1 to this number. Why? Because in opposite case we will have different numbers with the same rank, lying in the same colomn or row by properties of our connected component.
  5. Finally, we update ranks in table, we can use matrix A directly, because we update full connected component.

Complexity

Time complexity is O(mn *log(mn), because the most heavy part is sorting (in theory, in practice it works quite fast). Then even though we use dsu without optimizations, it will also work in average in logarithmic time, so total complexity will be the same. Space complexity is O(mn) to keep dictionary d.

Code

class DSU:
    def __init__(self, graph):
        self.p = {i:i for i in graph}

    def find(self, x):
        if self.p[x] != x:
            self.p[x] = self.find(self.p[x])
        return self.p[x]

    def union(self, x, y):
        self.p[self.find(x)] = self.find(y)
        
    def groups(self):
        ans = defaultdict(list)
        for el in self.p.keys():
            ans[self.find(el)].append(el)
        return ans
        
class Solution:
    def matrixRankTransform(self, M):
        n, m = len(M), len(M[0])
        rank = [0] * (m + n)
        d = defaultdict(list)
        
        for i, j in product(range(n), range(m)):
            d[M[i][j]].append([i, j])

        for a in sorted(d):
            graph = [i for i, j in d[a]] + [j + n for i, j in d[a]]
            dsu = DSU(graph)
            for i, j in d[a]: dsu.union(i, j + n)

            for group in dsu.groups().values():
                mx = max(rank[i] for i in group)
                for i in group: rank[i] = mx + 1
                    
            for i, j in d[a]: M[i][j] = rank[i]
            
        return M