1203. Sort Items by Groups Respecting Dependencies

link

It's easy to think that's a topological sort questions.

However, it used double topological sort (one for node, one for group)

And need to really care about the data structure, need to put group as hash mapping otherwise TLE.

class Solution:
    def sortItems(self, n: int, m: int, group: List[int], beforeItems: List[List[int]]) -> List[int]:

        indegN = [0] * n
        indegM = [0] * m

        ## indeg items relation
        M_L = defaultdict(list)
        ## node in group
        M_G = defaultdict(list)

        for i in range(n):
            M_G[group[i]].append(i)
            for ii in beforeItems[i]:
                M_L[ii].append(i)
                indegN[i] += 1
                # only card about the non -1 group relationship
                if group[i] != -1 and group[ii] != group[i]:
                    indegM[group[i]] += 1
        
        res = []
        # used items(for -1 group)
        S_N = set()
        dqM = deque([ i for i in range(m) if indegM[i] == 0 ])

        while True:
            dq1 = deque([i for i in M_G[-1] if i not in S_N and indegN[i] == 0])

            while dq1:
                idx = dq1.popleft()
                S_N.add(idx)
                res.append(idx)
                for ii in M_L[idx]:
                    indegN[ii] -= 1
                    if group[ii] != -1:
                        indegM[group[ii]] -= 1
                        if indegM[group[ii]] == 0:
                            dqM.append(group[ii])
                    if indegN[ii] == 0 and group[ii] == -1:
                        dq1.append(ii)

            # if can't find any group, just break the loop
            if not dqM:
                break

            while dqM:
                g = dqM.popleft()
                dq2 = deque([i for i in M_G[g] if indegN[i] == 0])

                while dq2:
                    idx = dq2.popleft()
                    res.append(idx)

                    for ii in M_L[idx]:
                        indegN[ii] -= 1
                        if group[ii] != -1 and group[idx] != group[ii]:
                            indegM[group[ii]] -= 1
                            if indegM[group[ii]] == 0:
                                dqM.append(group[ii])
                        if indegN[ii] == 0 and group[ii] == g:
                            dq2.append(ii)

        # might have unprocessed items
        if len(res) != n:
            return []

        return res