LC 1203 [H] Sort Items by Groups Respecting Dependencies - ALawliet/algorithms GitHub Wiki

topo sort the items and the groups and then look up the order of the items in the groups

class Solution:
    def sortItems(self, n: int, m: int, group: List[int], beforeItems: List[List[int]]) -> List[int]:
        # Helper function: returns topological order of a graph, if it exists.
        def get_top_order(graph, indegree):
            top_order = []
            Q = deque()
            for node in range(len(graph)):
                if indegree[node] == 0:
                    Q.append(node)
            while Q:
                v = Q.popleft()
                top_order.append(v)
                for u in graph[v]:
                    indegree[u] -= 1
                    if indegree[u] == 0:
                        Q.append(u)
            return top_order if len(top_order) == len(graph) else []

        # STEP 1: Create a new group for each item that belongs to no group. 
        for u in range(len(group)):
            if group[u] == -1:
                group[u] = m
                m+=1

        # STEP 2: Build directed graphs for items and groups.
        graph_items = [[] for _ in range(n)]
        indegree_items = [0] * n
        graph_groups = [[] for _ in range(m)]
        indegree_groups = [0] * m        
        for u in range(n):
            for v in beforeItems[u]:                
                graph_items[v].append(u)
                indegree_items[u] += 1
                if group[u]!=group[v]:
                    graph_groups[group[v]].append(group[u])
                    indegree_groups[group[u]] += 1                    

        # STEP 3: Find topological orders of items and groups.
        item_order = get_top_order(graph_items, indegree_items)
        group_order = get_top_order(graph_groups, indegree_groups)
        if not item_order or not group_order: return []

        # STEP 4: Find order of items within each group.
        order_within_group = collections.defaultdict(list)
        for v in item_order:
            order_within_group[group[v]].append(v)

        # STEP 5. Combine ordered groups.
        res = []
        for group in group_order:
            res += order_within_group[group]
        return res