Problem statement

https://binarysearch.com/problems/Matrix-Relations/

Solution

The idea is to use two topological sorts: one for columns and one for rows. What our topsort function will retrun is dictionary of levels: if we have graph with edges [0, 1], [1, 2], [0, 3], [3, 2], then we have on the first layer [0], then we have [1, 3] and finally we have [2]. It means that 0 must be on the first row, [1, 3] must be on the next row and finally 2 must be on the last row. Notice that it is given in problem statement that solution is unique. It means, that there will be always exactly n layers.

Complexity

It is O(V + E) for time and space.

Code

class Solution:
    def solve(self, N, relations):
        rdeg, cdeg = Counter(), Counter()
        rlist = defaultdict(list)
        clist = defaultdict(list)
        keys = set()

        for x, y, t in relations:
            keys |= set([x, y])
            xdeg = rdeg if t < 2 else cdeg
            xlist = rlist if t < 2 else clist
            a, b = [x, y][::1-(t%2)*2]
            xdeg[b] += 1
            xlist[a] += [b]
          
        def topsort(indeg, elist):
            ans = {}
            Q = deque([(x, 0) for x in keys if not indeg[x]])
    
            while Q:
                index, d = Q.popleft()
                ans[index] = d
                for v in elist[index]:
                    if indeg[v] > 0:
                        indeg[v] -= 1
                        if indeg[v] == 0:
                            Q.append((v, d + 1))
            return ans

        rans = topsort(rdeg, rlist)
        cans = topsort(cdeg, clist)
        ans = [[-1] * N for _ in range(N)]

        for x in keys:
            ans[cans[x]][rans[x]] = x
        return ans