Graphs

from collections import defaultdict
 
n = 4
edges = [(0, 1), (0, 2), (1, 2), (2, 3)]
 
adj: list[list[int]] = [[] for _ in range(n)]
for u, v in edges:  # undirected: both directions
    adj[u].append(v)
    adj[v].append(u)
print(adj)  # [[1, 2], [0, 2], [0, 1, 3], [2]]
 
nbrs: defaultdict[int, set[int]] = defaultdict(set)
for u, v in edges:
    nbrs[u].add(v)
    nbrs[v].add(u)
1 in nbrs[0]  # True: an O(1) edge test
 
matrix = [[0] * n for _ in range(n)]
for u, v in edges:
    matrix[u][v] = matrix[v][u] = 1
matrix[2]  # [1, 1, 0, 1]
 
roads = {"a": {"b": 4, "c": 1}, "c": {"b": 2}}
roads["a"]["c"]  # 1: the weight of a -> c

from collections import defaultdict
from collections.abc import Iterator
 
raw = "4 3\n1 2\n2 3\n3 1"  # n m, then m edges, 1-indexed
lines = raw.splitlines()
n, m = map(int, lines[0].split())
adj: list[list[int]] = [[] for _ in range(n + 1)]
for line in lines[1 : m + 1]:
    u, v = map(int, line.split())
    adj[u].append(v)
    adj[v].append(u)
print(adj[1:])
# [[2, 3], [1, 3], [2, 1], []]: node 4 has no edges
 
g: defaultdict[str, list[str]] = defaultdict(list)
g["a"].append("b")
"b" in g  # False: b has no entry of its own
g.get("b", [])  # []: reads without inserting
nodes = ["a", "b", "c"]
full = {u: g.get(u, []) for u in nodes}  # every node present
sorted(full)  # ['a', 'b', 'c']
 
 
def grid_neighbors(
    grid: list[str], r: int, c: int
) -> Iterator[tuple[int, int]]:
    for dr, dc in ((1, 0), (-1, 0), (0, 1), (0, -1)):
        nr, nc = r + dr, c + dc
        if 0 <= nr < len(grid) and 0 <= nc < len(grid[0]):
            if grid[nr][nc] != "#":
                yield nr, nc
 
 
list(grid_neighbors(["..#", ".#."], 0, 1))  # [(0, 0)]

from collections import deque
 
type Adj = list[list[int]]
 
 
def bfs(adj: Adj, src: int) -> list[int]:
    dist = [-1] * len(adj)  # -1: not reached yet
    dist[src] = 0
    q = deque([src])
    while q:
        u = q.popleft()
        for v in adj[u]:
            if dist[v] == -1:  # mark when enqueued
                dist[v] = dist[u] + 1
                q.append(v)
    return dist
 
 
def levels(adj: Adj, src: int) -> list[list[int]]:
    seen = {src}
    q = deque([src])
    out: list[list[int]] = []
    while q:
        level: list[int] = []
        for _ in range(len(q)):  # exactly this level
            u = q.popleft()
            level.append(u)
            for v in adj[u]:
                if v not in seen:
                    seen.add(v)
                    q.append(v)
        out.append(level)
    return out
 
 
adj: Adj = [[1, 2], [0, 3], [0, 3], [1, 2, 4], [3], []]
print(bfs(adj, 0))  # [0, 1, 1, 2, 3, -1]
print(levels(adj, 0))  # [[0], [1, 2], [3], [4]]

import math
from collections import deque
 
type WAdj = list[list[tuple[int, int]]]  # w is 0 or 1
 
 
def zero_one_bfs(adj: WAdj, src: int) -> list[float]:
    dist: list[float] = [math.inf] * len(adj)
    dist[src] = 0
    q = deque([src])
    while q:
        u = q.popleft()
        for v, w in adj[u]:
            if dist[u] + w < dist[v]:  # can improve later
                dist[v] = dist[u] + w
                if w == 0:
                    q.appendleft(v)  # same distance: front
                else:
                    q.append(v)
    return dist
 
 
wadj: WAdj = [[(1, 1), (2, 0)], [(3, 1)], [(1, 0)], []]
print(zero_one_bfs(wadj, 0))  # [0, 0, 0, 1]

type Adj = list[list[int]]
 
 
def count_reachable(adj: Adj, src: int) -> int:
    seen = [False] * len(adj)
 
    def dfs(u: int) -> None:  # recursive: fine below ~1,000
        seen[u] = True
        for v in adj[u]:
            if not seen[v]:
                dfs(v)
 
    dfs(src)
    return sum(seen)
 
 
def preorder(adj: Adj, src: int) -> list[int]:
    seen: set[int] = set()
    order: list[int] = []
    stack = [src]
    while stack:
        u = stack.pop()
        if u in seen:
            continue  # pushed twice; already visited
        seen.add(u)  # mark on pop: same order as recursion
        order.append(u)
        stack.extend(reversed(adj[u]))  # first child on top
    return order
 
 
def postorder(adj: Adj, src: int) -> list[int]:
    seen = {src}
    out: list[int] = []
    stack = [(src, iter(adj[src]))]
    while stack:
        u, it = stack[-1]
        for v in it:  # resumes where u left off
            if v not in seen:
                seen.add(v)
                stack.append((v, iter(adj[v])))
                break
        else:  # no unvisited neighbor left: u is finished
            stack.pop()
            out.append(u)
    return out
 
 
dag: Adj = [[1, 2], [3], [3], []]
print(count_reachable(dag, 0))  # 4
print(preorder(dag, 0))  # [0, 1, 3, 2]
print(postorder(dag, 0))  # [3, 1, 2, 0]
print(postorder(dag, 0)[::-1])  # [0, 2, 1, 3]: a topo order

import heapq
import math
 
type WAdj = list[list[tuple[int, int]]]  # (v, weight)
 
 
def dijkstra(
    adj: WAdj, src: int
) -> tuple[list[float], list[int]]:
    dist: list[float] = [math.inf] * len(adj)
    parent = [-1] * len(adj)
    dist[src] = 0
    h = [(0, src)]  # (distance so far, node)
    while h:
        d, u = heapq.heappop(h)
        if d > dist[u]:
            continue  # stale: u improved after this push
        for v, w in adj[u]:
            if d + w < dist[v]:
                dist[v] = d + w
                parent[v] = u
                heapq.heappush(h, (d + w, v))
    return dist, parent
 
 
def path_to(parent: list[int], t: int) -> list[int]:
    out: list[int] = []
    while t != -1:
        out.append(t)
        t = parent[t]
    return out[::-1]
 
 
adj: WAdj = [
    [(1, 4), (2, 1)],
    [(3, 1)],
    [(1, 2), (3, 5)],
    [],
]
dist, parent = dijkstra(adj, 0)
print(dist)  # [0, 3, 1, 4]
print(path_to(parent, 3))  # [0, 2, 1, 3]

import math
 
INF = math.inf
type WEdges = list[tuple[int, int, int]]  # (u, v, w)
 
 
def bellman_ford(
    n: int, edges: WEdges, src: int, rounds: int
) -> list[float]:
    dist: list[float] = [INF] * n
    dist[src] = 0
    for _ in range(rounds):  # n - 1 rounds: all paths
        nxt = dist[:]  # read last round's values only
        for u, v, w in edges:
            if dist[u] + w < nxt[v]:
                nxt[v] = dist[u] + w
        dist = nxt
    return dist
 
 
flights: WEdges = [(0, 1, 100), (1, 2, 100), (0, 2, 500)]
bellman_ford(3, flights, 0, 3 - 1)[2]  # 200
bellman_ford(3, flights, 0, 1)[2]  # 500: one edge at most
 
n = 3
d: list[list[float]] = [[INF] * n for _ in range(n)]
for i in range(n):
    d[i][i] = 0
for u, v, w in [(0, 1, 4), (1, 2, -2), (0, 2, 5)]:
    d[u][v] = min(d[u][v], w)
for k in range(n):  # Floyd–Warshall: k is the outer loop
    for i in range(n):
        for j in range(n):
            if d[i][k] + d[k][j] < d[i][j]:
                d[i][j] = d[i][k] + d[k][j]
print(d[0])  # [0, 4, 2]

from graphlib import CycleError, TopologicalSorter
 
deps = {"deploy": {"test", "build"}, "test": {"build"}}
print(list(TopologicalSorter(deps).static_order()))
# ['build', 'test', 'deploy']
 
cyclic = TopologicalSorter({"a": {"b"}, "b": {"a"}})
try:
    list(cyclic.static_order())
except CycleError as e:
    print(e.args[1])  # ['a', 'b', 'a']
 
ts: TopologicalSorter[str] = TopologicalSorter()
ts.add("c", "a", "b")  # c needs a and b
ts.add("d", "c")
ts.prepare()  # raises CycleError here if there is one
batches: list[list[str]] = []
while ts.is_active():
    ready = ts.get_ready()  # everything runnable now
    batches.append(sorted(ready))
    ts.done(*ready)
print(batches)  # [['a', 'b'], ['c'], ['d']]

from collections import deque
 
type Edges = list[tuple[int, int]]  # (u, v): u before v
 
 
def kahn(n: int, edges: Edges) -> list[int] | None:
    adj: list[list[int]] = [[] for _ in range(n)]
    indeg = [0] * n
    for u, v in edges:
        adj[u].append(v)
        indeg[v] += 1
    q = deque(u for u in range(n) if indeg[u] == 0)
    order: list[int] = []
    while q:
        u = q.popleft()
        order.append(u)
        for v in adj[u]:
            indeg[v] -= 1
            if indeg[v] == 0:
                q.append(v)
    if len(order) < n:
        return None  # some nodes never freed: a cycle
    return order
 
 
diamond = [(0, 1), (0, 2), (1, 3), (2, 3)]
print(kahn(4, diamond))  # [0, 1, 2, 3]
print(kahn(2, [(0, 1), (1, 0)]))  # None

class DSU:
    def __init__(self, n: int) -> None:
        self.parent = list(range(n))
        self.size = [1] * n
        self.count = n  # number of components
 
    def find(self, x: int) -> int:
        p = self.parent
        while p[x] != x:
            p[x] = p[p[x]]  # path halving
            x = p[x]
        return x
 
    def union(self, a: int, b: int) -> bool:
        ra, rb = self.find(a), self.find(b)
        if ra == rb:
            return False  # same group: this edge is a cycle
        if self.size[ra] < self.size[rb]:
            ra, rb = rb, ra  # hang the smaller tree
        self.parent[rb] = ra
        self.size[ra] += self.size[rb]
        self.count -= 1
        return True
 
 
dsu = DSU(5)
dsu.union(0, 1), dsu.union(1, 2), dsu.union(0, 2)
# (True, True, False): the third edge closes a cycle
dsu.count  # 3: {0, 1, 2}, {3}, {4}
 
parent: dict[str, str] = {}  # string or sparse nodes
 
 
def find(x: str) -> str:
    parent.setdefault(x, x)  # first sight: its own root
    while parent[x] != x:
        parent[x] = parent[parent[x]]
        x = parent[x]
    return x
 
 
parent[find("ann")] = find("bo")  # union
find("ann") == find("bo"), find("cy") == find("bo")
# (True, False)

from functools import cache
 
type Adj = list[list[int]]
 
 
def all_paths(adj: Adj, target: int) -> list[list[int]]:
    out: list[list[int]] = []
    path = [0]
 
    def go(u: int) -> None:
        if u == target:
            out.append(path[:])  # copy: path keeps changing
            return
        for v in adj[u]:
            path.append(v)
            go(v)
            path.pop()
 
    go(0)
    return out
 
 
def count_paths(adj: Adj, target: int) -> int:
    @cache
    def ways(u: int) -> int:  # a DAG, so no cycles
        if u == target:
            return 1
        return sum(ways(v) for v in adj[u])
 
    return ways(0)
 
 
dag: Adj = [[1, 2], [3], [3], []]
print(all_paths(dag, 3))  # [[0, 1, 3], [0, 2, 3]]
print(count_paths(dag, 3))  # 2

from collections import deque
from collections.abc import Hashable, Iterable, Mapping
 
type Graph = dict[str, list[str]]
type Adj = list[list[int]]
 
 
def bfs_order[T: Hashable](
    g: Mapping[T, Iterable[T]], start: T
) -> list[T]:
    seen = {start}
    order: list[T] = []
    q = deque([start])
    while q:
        u = q.popleft()
        order.append(u)
        for v in g.get(u, ()):
            if v not in seen:
                seen.add(v)
                q.append(v)
    return order
 
 
g: Graph = {"a": ["b", "c"], "b": ["d"]}
print(bfs_order(g, "a"))  # ['a', 'b', 'c', 'd']
adj: Adj = [[1], [0, 2], []]
print(bfs_order(dict(enumerate(adj)), 0))  # [0, 1, 2]

from graphlib import CycleError, TopologicalSorter
 
 
def find_order(
    n: int, prereqs: list[list[int]]
) -> list[int]:
    ts: TopologicalSorter[int] = TopologicalSorter()
    for course in range(n):
        ts.add(course)  # courses without prerequisites too
    for course, before in prereqs:
        ts.add(course, before)  # before comes first
    try:
        return list(ts.static_order())
    except CycleError:
        return []
 
 
reqs = [[1, 0], [2, 0], [3, 1], [3, 2]]
order = find_order(4, reqs)
pos = {c: i for i, c in enumerate(order)}
print(all(pos[b] < pos[a] for a, b in reqs))  # True
print(order[0], order[-1])  # 0 3
print(find_order(2, [[1, 0], [0, 1]]))  # []

class Node:
    def __init__(self, val: int) -> None:
        self.val = val
        self.neighbors: list[Node] = []
 
 
def clone(start: Node | None) -> Node | None:
    if start is None:
        return None
    copies = {start: Node(start.val)}  # old -> new
    stack = [start]
    while stack:
        old = stack.pop()
        for nb in old.neighbors:
            if nb not in copies:  # create on first sight
                copies[nb] = Node(nb.val)
                stack.append(nb)
            copies[old].neighbors.append(copies[nb])
    return copies[start]
 
 
a, b, c = Node(1), Node(2), Node(3)
a.neighbors = [b, c]
b.neighbors = [a, c]
c.neighbors = [a, b]
a2 = clone(a)
assert a2 is not None
print(a2.val, [n.val for n in a2.neighbors])  # 1 [2, 3]
print(a2 is not a, a2.neighbors[0].neighbors[0] is a2)
# True True: a new graph with the same shape

from collections import defaultdict
 
 
def merge(accounts: list[list[str]]) -> list[list[str]]:
    parent: dict[str, str] = {}
 
    def find(x: str) -> str:
        parent.setdefault(x, x)
        while parent[x] != x:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x
 
    owner: dict[str, str] = {}
    for name, first, *rest in accounts:
        owner[first] = name
        find(first)  # register one-email accounts too
        for email in rest:
            owner[email] = name
            parent[find(email)] = find(first)
    groups: defaultdict[str, list[str]] = defaultdict(list)
    for email in owner:
        groups[find(email)].append(email)
    return sorted(
        [owner[root], *sorted(emails)]
        for root, emails in groups.items()
    )
 
 
accounts = [
    ["ann", "a@x", "b@x"],
    ["ann", "a@x", "c@x"],
    ["bo", "d@x"],
    ["ann", "e@x"],
]
for row in merge(accounts):
    print(row)
# ['ann', 'a@x', 'b@x', 'c@x']
# ['ann', 'e@x']
# ['bo', 'd@x']

WHITE, GRAY, BLACK = 0, 1, 2  # unseen, on the path, done
 
 
def has_cycle(adj: list[list[int]]) -> bool:
    color = [WHITE] * len(adj)
    for s in range(len(adj)):
        if color[s] != WHITE:
            continue
        color[s] = GRAY
        stack = [(s, iter(adj[s]))]
        while stack:
            u, it = stack[-1]
            for v in it:
                if color[v] == GRAY:
                    return True  # back edge to the path
                if color[v] == WHITE:
                    color[v] = GRAY
                    stack.append((v, iter(adj[v])))
                    break
            else:
                color[u] = BLACK  # all neighbors finished
                stack.pop()
    return False
 
 
print(has_cycle([[1], [2], [0]]))  # True
print(has_cycle([[1, 2], [3], [3], []]))  # False: a diamond
chain = [[i + 1] for i in range(99_999)] + [[]]
print(has_cycle(chain))  # False: 100,000 deep, no recursion

from collections import defaultdict, deque
 
 
def calc(
    eqs: list[tuple[str, str]],
    vals: list[float],
    queries: list[tuple[str, str]],
) -> list[float]:
    g: defaultdict[str, dict[str, float]] = defaultdict(dict)
    for (a, b), k in zip(eqs, vals, strict=True):
        g[a][b] = k  # a / b = k
        g[b][a] = 1 / k
    out: list[float] = []
    for src, dst in queries:
        if src not in g or dst not in g:  # in: no insert
            out.append(-1.0)
            continue
        ratio = {src: 1.0}  # src / node, per reached node
        q = deque([src])
        while q:
            u = q.popleft()
            for v, k in g[u].items():
                if v not in ratio:
                    ratio[v] = ratio[u] * k
                    q.append(v)
        out.append(ratio.get(dst, -1.0))
    return out
 
 
eqs = [("a", "b"), ("b", "c")]
qs = [("a", "c"), ("b", "a"), ("a", "e"), ("x", "x")]
print(calc(eqs, [2.0, 3.0], qs))  # [6.0, 0.5, -1.0, -1.0]

import heapq
 
 
def alien_order(words: list[str]) -> str:
    chars = {c for w in words for c in w}
    adj: dict[str, set[str]] = {c: set() for c in chars}
    indeg = dict.fromkeys(chars, 0)
    for a, b in zip(words, words[1:]):
        for x, y in zip(a, b):
            if x != y:  # the first difference decides
                if y not in adj[x]:
                    adj[x].add(y)
                    indeg[y] += 1
                break
        else:
            if len(a) > len(b):
                return ""  # "abc" before "ab": invalid
    h = [c for c in chars if indeg[c] == 0]
    heapq.heapify(h)  # smallest ready letter first
    out: list[str] = []
    while h:
        c = heapq.heappop(h)
        out.append(c)
        for d in adj[c]:
            indeg[d] -= 1
            if indeg[d] == 0:
                heapq.heappush(h, d)
    return "".join(out) if len(out) == len(chars) else ""
 
 
words = ["wrt", "wrf", "er", "ett", "rftt"]
print(alien_order(words))  # wertf
print(alien_order(["ca", "cb"]))  # abc: c waits its turn
print(repr(alien_order(["z", "x", "z"])))  # '': a cycle