"""Helpers for manipulations with graphs."""

from __future__ import annotations

from collections.abc import Iterator, Set as AbstractSet

from typing import TypeVar

T = TypeVar("T")

def strongly_connected_components(

vertices: AbstractSet[T], edges: dict[T, list[T]]

) -> Iterator[set[T]]:

"""Compute Strongly Connected Components of a directed graph.

Args:

vertices: the labels for the vertices

edges: for each vertex, gives the target vertices of its outgoing edges

Returns:

An iterator yielding strongly connected components, each

represented as a set of vertices. Each input vertex will occur

exactly once; vertices not part of a SCC are returned as

singleton sets.

From https://code.activestate.com/recipes/578507/.

"""

identified: set[T] = set()

stack: list[T] = []

index: dict[T, int] = {}

boundaries: list[int] = []

def dfs(v: T) -> Iterator[set[T]]:

index[v] = len(stack)

stack.append(v)

boundaries.append(index[v])

for w in edges[v]:

if w not in index:

yield from dfs(w)

elif w not in identified:

while index[w] < boundaries[-1]:

boundaries.pop()

if boundaries[-1] == index[v]:

boundaries.pop()

scc = set(stack[index[v] :])

del stack[index[v] :]

identified.update(scc)

yield scc

for v in vertices:

if v not in index:

yield from dfs(v)

def prepare_sccs(

sccs: list[set[T]], edges: dict[T, list[T]]

) -> dict[AbstractSet[T], set[AbstractSet[T]]]:

"""Use original edges to organize SCCs in a graph by dependencies between them."""

sccsmap = {}

for scc in sccs:

scc_frozen = frozenset(scc)

for v in scc:

sccsmap[v] = scc_frozen

data: dict[AbstractSet[T], set[AbstractSet[T]]] = {}

for scc in sccs:

deps: set[AbstractSet[T]] = set()

for v in scc:

deps.update(sccsmap[x] for x in edges[v])

data[frozenset(scc)] = deps

return data

class topsort(Iterator[set[T]]): # noqa: N801

"""Topological sort using Kahn's algorithm.

Uses in-degree counters and a reverse adjacency list, so the total work

is O(V + E).

Implemented as a class rather than a generator for better mypyc

compilation.

Args:

data: A map from vertices to all vertices that it has an edge

connecting it to. NOTE: dependency sets in this data

structure are modified in place to remove self-dependencies.

Orphans are handled internally and are not added to `data`.

Returns:

An iterator yielding sets of vertices that have an equivalent

ordering.

Example:

Suppose the input has the following structure:

{A: {B, C}, B: {D}, C: {D}}

The algorithm treats orphan dependencies as if normalized to:

{A: {B, C}, B: {D}, C: {D}, D: {}}

It will yield the following values:

{D}

{B, C}

{A}

"""

def __init__(self, data: dict[T, set[T]]) -> None:

# Single pass: remove self-deps, build reverse adjacency list,

# compute in-degree counts, detect orphans, and find initial ready set.

in_degree: dict[T, int] = {}

rev: dict[T, list[T]] = {}

ready: set[T] = set()

for item, deps in data.items():

deps.discard(item) # Ignore self dependencies.

deg = len(deps)

in_degree[item] = deg

if deg == 0:

ready.add(item)

if item not in rev:

rev[item] = []

for dep in deps:

if dep in rev:

rev[dep].append(item)

else:

rev[dep] = [item]

if dep not in data:

# Orphan: appears as dependency but has no entry in data.

in_degree[dep] = 0

ready.add(dep)

self.in_degree = in_degree

self.rev = rev

self.ready = ready

self.remaining = len(in_degree) - len(ready)

def __iter__(self) -> Iterator[set[T]]:

return self

def __next__(self) -> set[T]:

ready = self.ready

if not ready:

assert self.remaining == 0, (

f"A cyclic dependency exists amongst "

f"{[k for k, deg in self.in_degree.items() if deg > 0]!r}"

)

raise StopIteration

in_degree = self.in_degree

rev = self.rev

new_ready: set[T] = set()

for item in ready:

for dependent in rev[item]:

new_deg = in_degree[dependent] - 1

in_degree[dependent] = new_deg

if new_deg == 0:

new_ready.add(dependent)

self.remaining -= len(new_ready)

self.ready = new_ready

return ready