Leetcode 1257 Solution

class Solution:
    def matrixRankTransform(self, matrix: List[List[int]]) -> List[List[int]]:
        m = len(matrix)
        n = len(matrix[0])

        row_ranks = [0] * m
        col_ranks = [0] * m

        nums = defaultdict(list)

        for i in range(len(matrix)):
            for j in range(len(matrix[i])):
                num = matrix[i][j]
                nums[num].append((i, j))

        res = [[0] * n for _ in range(m)]

        f = [0] * 1004
        ranks = [0] * 1004

        def find(i):
            nonlocal f

            if f[i] == i:
                return i

            f[i] = find(f[i])
            return f[i]

        for num in sorted(nums.keys()):
            positions = nums[num]

            for i, j in positions:
                f[i] = i
                f[j + m] = j + m

                ranks[i] = row_ranks[i]
                ranks[j + m] = col_ranks[j]

            for i, j in positions:
                root_i = find(i)
                root_j = find(j + m)

                if root_i != root_j:
                    f[root_i] = root_j
                    ranks[root_j] = max(ranks[root_i], ranks[root_j])

            for i, j in positions:
                root_i = find(i)
                rank = ranks[root_i] + 1
                res[i][j] = rank
                row_ranks[i] = rank
                col_ranks[j] = rank


        return res