Leetcode 1203. Sort Items by Groups Respecting Dependencies

Problem statement

https://leetcode.com/problems/sort-items-by-groups-respecting-dependencies/

Solution

The idea is to do several topological sorts: first one for groups, then inside each group. Here I use template for topological sort with use of bfs, because it is shorter in this way. Then for each group with one element (that is with group = -1), we put it to separate group with one element. Then we iterato over all before elements and decide whith group to update: group_graphs: defaultdict of list inside groups and between_groups for elements from different groups.

Complexity

Time complexity is O(n + k), where n is number of items and k is number of edges in total.

Code

class Solution:
    def sortItems(self, n, m, group, before):
        def Topological(verts, edges):
            pre, suc = defaultdict(int), defaultdict(list)
            for a, b in edges:
                pre[a] += 1
                suc[b].append(a)

            free, out = set(verts) - set(pre), []
            while free:
                a = free.pop()
                out.append(a)
                for b in suc[a]:
                    pre[b] -= 1
                    if not pre[b]: free.add(b)
            return out * (len(out) == len(verts))
        
        t, ans = max(group) + 1, []
        for i, g in enumerate(group):
            if g == -1:
                group[i] = t
                t += 1
        
        G = defaultdict(list)
        for i, g in zip(range(n), group):
            G[g].append(i)
                
        group_graphs = defaultdict(list)
        between_groups = []
        
        for i in range(n):
            for j in before[i]:
                if group[i] == group[j]:
                    group_graphs[group[i]].append((i, j))
                else:
                    between_groups.append((group[i], group[j]))
                    
        q = Topological(list(range(t)), between_groups)
        
        if not q: return []
        for i in range(t):
            ans.extend(Topological(G[q[i]], group_graphs[q[i]]))
    
        return ans if len(ans) == n else []