cp-library
This documentation is automatically generated by online-judge-tools/verification-helper
test/atcoder/abc/abc218_f_shortest_path_weighted.test.py
- View this file on GitHub
- Last update: 2025-03-30 20:17:47+09:00
- Problem: https://atcoder.jp/contests/abc218/tasks/abc218_f
Depends on
- cp_library/alg/graph/dfs_options_cls.py
- cp_library/alg/graph/digraph_cls.py
- cp_library/alg/graph/digraph_weighted_cls.py
- cp_library/alg/graph/edge_cls.py
- cp_library/alg/graph/edge_list_cls.py
- cp_library/alg/graph/edge_weighted_cls.py
- cp_library/alg/graph/graph_proto.py
- cp_library/alg/graph/graph_weighted_proto.py
- cp_library/ds/csr/csr_incremental_cls.py
- cp_library/ds/dsu_cls.py
- cp_library/ds/elist_fn.py
- cp_library/ds/heap/heap_proto.py
- cp_library/ds/heap/priority_queue_cls.py
- cp_library/io/fast_io_cls.py
- cp_library/io/parser_cls.py
- cp_library/io/read_fn.py
- cp_library/io/write_fn.py
Code
# verification-helper: PROBLEM https://atcoder.jp/contests/abc218/tasks/abc218_f
from math import inf
def main():
N, M = read(tuple[int, ...])
E = read(EdgeList[M])
EW = [(u,v,1) for u,v in E]
G = DiGraphWeighted(N,EW)
path = G.shortest_path(0,N-1)
if path is None:
shortest = -1
else:
path = set(path)
shortest = len(path)
for e in range(M):
if path is not None and e in path:
G2 = DiGraph(N, E[:e]+E[e+1:])
ans = G2.distance(0,N-1)
else:
ans = shortest
write(ans if ans != inf else -1)
from cp_library.alg.graph.edge_list_cls import EdgeList
from cp_library.alg.graph.digraph_weighted_cls import DiGraphWeighted
from cp_library.alg.graph.digraph_cls import DiGraph
from cp_library.io.read_fn import read
from cp_library.io.write_fn import write
if __name__ == "__main__":
main()# verification-helper: PROBLEM https://atcoder.jp/contests/abc218/tasks/abc218_f
from math import inf
def main():
N, M = read(tuple[int, ...])
E = read(EdgeList[M])
EW = [(u,v,1) for u,v in E]
G = DiGraphWeighted(N,EW)
path = G.shortest_path(0,N-1)
if path is None:
shortest = -1
else:
path = set(path)
shortest = len(path)
for e in range(M):
if path is not None and e in path:
G2 = DiGraph(N, E[:e]+E[e+1:])
ans = G2.distance(0,N-1)
else:
ans = shortest
write(ans if ans != inf else -1)
'''
╺━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╸
https://kobejean.github.io/cp-library
'''
from typing import TypeVar
import typing
from collections import deque
from numbers import Number
from types import GenericAlias
from typing import Callable, Collection, Iterator, Union
import os
import sys
from io import BytesIO, IOBase
class FastIO(IOBase):
BUFSIZE = 8192
newlines = 0
def __init__(self, file):
self._fd = file.fileno()
self.buffer = BytesIO()
self.writable = "x" in file.mode or "r" not in file.mode
self.write = self.buffer.write if self.writable else None
def read(self):
BUFSIZE = self.BUFSIZE
while True:
b = os.read(self._fd, max(os.fstat(self._fd).st_size, BUFSIZE))
if not b:
break
ptr = self.buffer.tell()
self.buffer.seek(0, 2), self.buffer.write(b), self.buffer.seek(ptr)
self.newlines = 0
return self.buffer.read()
def readline(self):
BUFSIZE = self.BUFSIZE
while self.newlines == 0:
b = os.read(self._fd, max(os.fstat(self._fd).st_size, BUFSIZE))
self.newlines = b.count(b"\n") + (not b)
ptr = self.buffer.tell()
self.buffer.seek(0, 2), self.buffer.write(b), self.buffer.seek(ptr)
self.newlines -= 1
return self.buffer.readline()
def flush(self):
if self.writable:
os.write(self._fd, self.buffer.getvalue())
self.buffer.truncate(0), self.buffer.seek(0)
class IOWrapper(IOBase):
stdin: 'IOWrapper' = None
stdout: 'IOWrapper' = None
def __init__(self, file):
self.buffer = FastIO(file)
self.flush = self.buffer.flush
self.writable = self.buffer.writable
def write(self, s):
return self.buffer.write(s.encode("ascii"))
def read(self):
return self.buffer.read().decode("ascii")
def readline(self):
return self.buffer.readline().decode("ascii")
sys.stdin = IOWrapper.stdin = IOWrapper(sys.stdin)
sys.stdout = IOWrapper.stdout = IOWrapper(sys.stdout)
_T = TypeVar('T')
class TokenStream(Iterator):
stream = IOWrapper.stdin
def __init__(self):
self.queue = deque()
def __next__(self):
if not self.queue: self.queue.extend(self._line())
return self.queue.popleft()
def wait(self):
if not self.queue: self.queue.extend(self._line())
while self.queue: yield
def _line(self):
return TokenStream.stream.readline().split()
def line(self):
if self.queue:
A = list(self.queue)
self.queue.clear()
return A
return self._line()
TokenStream.default = TokenStream()
class CharStream(TokenStream):
def _line(self):
return TokenStream.stream.readline().rstrip()
CharStream.default = CharStream()
ParseFn = Callable[[TokenStream],_T]
class Parser:
def __init__(self, spec: Union[type[_T],_T]):
self.parse = Parser.compile(spec)
def __call__(self, ts: TokenStream) -> _T:
return self.parse(ts)
@staticmethod
def compile_type(cls: type[_T], args = ()) -> _T:
if issubclass(cls, Parsable):
return cls.compile(*args)
elif issubclass(cls, (Number, str)):
def parse(ts: TokenStream): return cls(next(ts))
return parse
elif issubclass(cls, tuple):
return Parser.compile_tuple(cls, args)
elif issubclass(cls, Collection):
return Parser.compile_collection(cls, args)
elif callable(cls):
def parse(ts: TokenStream):
return cls(next(ts))
return parse
else:
raise NotImplementedError()
@staticmethod
def compile(spec: Union[type[_T],_T]=int) -> ParseFn[_T]:
if isinstance(spec, (type, GenericAlias)):
cls = typing.get_origin(spec) or spec
args = typing.get_args(spec) or tuple()
return Parser.compile_type(cls, args)
elif isinstance(offset := spec, Number):
cls = type(spec)
def parse(ts: TokenStream): return cls(next(ts)) + offset
return parse
elif isinstance(args := spec, tuple):
return Parser.compile_tuple(type(spec), args)
elif isinstance(args := spec, Collection):
return Parser.compile_collection(type(spec), args)
elif isinstance(fn := spec, Callable):
def parse(ts: TokenStream): return fn(next(ts))
return parse
else:
raise NotImplementedError()
@staticmethod
def compile_line(cls: _T, spec=int) -> ParseFn[_T]:
if spec is int:
fn = Parser.compile(spec)
def parse(ts: TokenStream): return cls([int(token) for token in ts.line()])
return parse
else:
fn = Parser.compile(spec)
def parse(ts: TokenStream): return cls([fn(ts) for _ in ts.wait()])
return parse
@staticmethod
def compile_repeat(cls: _T, spec, N) -> ParseFn[_T]:
fn = Parser.compile(spec)
def parse(ts: TokenStream): return cls([fn(ts) for _ in range(N)])
return parse
@staticmethod
def compile_children(cls: _T, specs) -> ParseFn[_T]:
fns = tuple((Parser.compile(spec) for spec in specs))
def parse(ts: TokenStream): return cls([fn(ts) for fn in fns])
return parse
@staticmethod
def compile_tuple(cls: type[_T], specs) -> ParseFn[_T]:
if isinstance(specs, (tuple,list)) and len(specs) == 2 and specs[1] is ...:
return Parser.compile_line(cls, specs[0])
else:
return Parser.compile_children(cls, specs)
@staticmethod
def compile_collection(cls, specs):
if not specs or len(specs) == 1 or isinstance(specs, set):
return Parser.compile_line(cls, *specs)
elif (isinstance(specs, (tuple,list)) and len(specs) == 2 and isinstance(specs[1], int)):
return Parser.compile_repeat(cls, specs[0], specs[1])
else:
raise NotImplementedError()
class Parsable:
@classmethod
def compile(cls):
def parser(ts: TokenStream): return cls(next(ts))
return parser
class Edge(tuple, Parsable):
@classmethod
def compile(cls, I=-1):
def parse(ts: TokenStream):
u,v = ts.line()
return cls((int(u)+I,int(v)+I))
return parse
E = TypeVar('E', bound=Edge)
M = TypeVar('M', bound=int)
class EdgeCollection(Parsable):
@classmethod
def compile(cls, M: M, E: E = Edge[-1]):
if isinstance(I := E, int):
E = Edge[I]
edge = Parser.compile(E)
def parse(ts: TokenStream):
return cls(edge(ts) for _ in range(M))
return parse
class EdgeList(EdgeCollection, list[E]):
pass
class EdgeSet(EdgeCollection, set[E]):
pass
from functools import total_ordering
@total_ordering
class EdgeWeighted(Edge):
def __lt__(self, other: tuple) -> bool:
a = self[2],self[0],self[1]
b = other[2],other[0],other[1]
return a < b
@classmethod
def compile(cls, I=-1):
def parse(ts: TokenStream):
u,v,w = ts.line()
return cls((int(u)+I, int(v)+I, int(w)))
return parse
from typing import Union, overload
from heapq import heapify, heappop, heappush
import operator
from enum import auto, IntFlag, IntEnum
class DFSFlags(IntFlag):
ENTER = auto()
DOWN = auto()
BACK = auto()
CROSS = auto()
LEAVE = auto()
UP = auto()
MAXDEPTH = auto()
RETURN_PARENTS = auto()
RETURN_DEPTHS = auto()
BACKTRACK = auto()
CONNECT_ROOTS = auto()
# Common combinations
ALL_EDGES = DOWN | BACK | CROSS
EULER_TOUR = DOWN | UP
INTERVAL = ENTER | LEAVE
TOPDOWN = DOWN | CONNECT_ROOTS
BOTTOMUP = UP | CONNECT_ROOTS
RETURN_ALL = RETURN_PARENTS | RETURN_DEPTHS
class DFSEvent(IntEnum):
ENTER = DFSFlags.ENTER
DOWN = DFSFlags.DOWN
BACK = DFSFlags.BACK
CROSS = DFSFlags.CROSS
LEAVE = DFSFlags.LEAVE
UP = DFSFlags.UP
MAXDEPTH = DFSFlags.MAXDEPTH
def elist(est_len: int) -> list: ...
try:
from __pypy__ import newlist_hint
except:
def newlist_hint(hint):
return []
elist = newlist_hint
from typing import Iterable, Union, overload
class GraphProtocol(list, Parsable):
def __init__(G, N: int, E: list = None, adj: Iterable = None):
G.N = N
if E is not None:
G.M, G.E = len(E), E
if adj is not None:
super().__init__(adj)
def neighbors(G, v: int) -> Iterable[int]:
return G[v]
def edge_ids(G) -> list[list[int]]: ...
@overload
def distance(G) -> list[list[int]]: ...
@overload
def distance(G, s: int = 0) -> list[int]: ...
@overload
def distance(G, s: int, g: int) -> int: ...
def distance(G, s = None, g = None):
if s == None:
return G.floyd_warshall()
else:
return G.bfs(s, g)
@overload
def bfs(G, s: Union[int,list] = 0) -> list[int]: ...
@overload
def bfs(G, s: Union[int,list], g: int) -> int: ...
def bfs(G, s = 0, g = None):
D = [inf for _ in range(G.N)]
q = deque([s] if isinstance(s, int) else s)
for u in q: D[u] = 0
while q:
nd = D[u := q.popleft()]+1
if u == g: return D[u]
for v in G.neighbors(u):
if nd < D[v]:
D[v] = nd
q.append(v)
return D if g is None else inf
@overload
def shortest_path(G, s: int, g: int) -> Union[list[int],None]: ...
@overload
def shortest_path(G, s: int, g: int, distances = True) -> tuple[Union[list[int],None],list[int]]: ...
def shortest_path(G, s: int, g: int, distances = False) -> list[int]:
D = [inf] * G.N
D[s] = 0
if s == g:
return ([], D) if distances else []
par = [-1] * G.N
par_edge = [-1] * G.N
Eid = G.edge_ids()
q = deque([s])
while q:
nd = D[u := q.popleft()] + 1
if u == g: break
for v, eid in zip(G[u], Eid[u]):
if nd < D[v]:
D[v] = nd
par[v] = u
par_edge[v] = eid
q.append(v)
if D[g] == inf:
return (None, D) if distances else None
path = []
current = g
while current != s:
path.append(par_edge[current])
current = par[current]
return (path[::-1], D) if distances else path[::-1]
def floyd_warshall(G) -> list[list[int]]:
D = [[inf]*G.N for _ in range(G.N)]
for u in range(G.N):
D[u][u] = 0
for v in G.neighbors(u):
D[u][v] = 1
for k, Dk in enumerate(D):
for Di in D:
if Di[k] == inf: continue
for j in range(G.N):
if Dk[j] == inf: continue
Di[j] = min(Di[j], Di[k]+Dk[j])
return D
def find_cycle(G, s = 0, vis = None, par = None):
N = G.N
vis = vis or [0] * N
par = par or [-1] * N
if vis[s]: return None
vis[s] = 1
stack = [(True, s)]
while stack:
forw, v = stack.pop()
if forw:
stack.append((False, v))
vis[v] = 1
for u in G.neighbors(v):
if vis[u] == 1 and u != par[v]:
# Cycle detected
cyc = [u]
vis[u] = 2
while v != u:
cyc.append(v)
vis[v] = 2
v = par[v]
return cyc
elif vis[u] == 0:
par[u] = v
stack.append((True, u))
else:
vis[v] = 2
return None
def find_minimal_cycle(G, s=0):
D, par, que = [inf] * (N := G.N), [-1] * N, deque([s])
D[s] = 0
while que:
for v in G[u := que.popleft()]:
if v == s: # Found cycle back to start
cycle = [u]
while u != s: cycle.append(u := par[u])
return cycle
if D[v] < inf: continue
D[v], par[v] = D[u]+1, u
que.append(v)
def bridges(G):
tin = [-1] * G.N
low = [-1] * G.N
par = [-1] * G.N
vis = [0] * G.N
in_edge = [-1] * G.N
Eid = G.edge_ids()
time = 0
bridges = []
stack = list(range(G.N))
while stack:
p = par[v := stack.pop()]
if vis[v] == 0:
vis[v] = 1
tin[v] = low[v] = time
time += 1
stack.append(v)
for i, child in enumerate(G.neighbors(v)):
if child == p: continue
if vis[child] == 0: # Tree edge - recurse
par[child] = v
in_edge[child] = Eid[v][i]
stack.append(child)
else: # Back edge - update low-link value
low[v] = min(low[v], tin[child])
elif vis[v] == 1:
vis[v] = 2
if p != -1:
low[p] = min(low[p], low[v])
if low[v] > tin[p]: bridges.append(in_edge[v])
return bridges
def articulation_points(G):
'''
Find articulation points in an undirected graph using DFS events.
Returns a boolean list that is True for indices where the vertex is an articulation point.
'''
N = G.N
order = [-1] * N
low = [-1] * N
par = [-1] * N
state = [0] * N
children = [0] * N
ap = [False] * N
time = 0
stack = list(range(N))
while stack:
v = stack.pop()
p = par[v]
if state[v] == 0:
state[v] = 1
order[v] = low[v] = time
time += 1
stack.append(v)
for child in G[v]:
if order[child] == -1:
par[child] = v
stack.append(child)
elif child != p:
low[v] = min(low[v], order[child])
if p != -1:
children[p] += 1
elif state[v] == 1:
state[v] = 2
ap[v] |= p == -1 and children[v] > 1
if p != -1:
low[p] = min(low[p], low[v])
ap[p] |= par[p] != -1 and low[v] >= order[p]
return ap
def dfs_events(G, flags: DFSFlags, s: Union[int,list,None] = None, max_depth: Union[int,None] = None):
if flags == DFSFlags.INTERVAL:
if max_depth is None:
return G.dfs_enter_leave(s)
elif flags == DFSFlags.DOWN or flags == DFSFlags.TOPDOWN:
if max_depth is None:
edges = G.dfs_topdown(s, DFSFlags.CONNECT_ROOTS in flags)
return [(DFSEvent.DOWN, p, u) for p,u in edges]
elif flags == DFSFlags.UP or flags == DFSFlags.BOTTOMUP:
if max_depth is None:
edges = G.dfs_bottomup(s, DFSFlags.CONNECT_ROOTS in flags)
return [(DFSEvent.UP, p, u) for p,u in edges]
elif flags & DFSFlags.BACKTRACK:
return G.dfs_backtrack(flags, s, max_depth)
state = [0] * G.N
child = [0] * G.N
stack = [0] * G.N
if flags & DFSFlags.RETURN_PARENTS:
parents = [-1] * G.N
if flags & DFSFlags.RETURN_DEPTHS:
depths = [-1] * G.N
events = []
for s in G.starts(s):
stack[depth := 0] = s
if (DFSFlags.DOWN|DFSFlags.CONNECT_ROOTS) in flags:
events.append((DFSEvent.DOWN,-1,s))
while depth != -1:
u = stack[depth]
if not state[u]:
state[u] = 1
if flags & DFSFlags.ENTER:
events.append((DFSEvent.ENTER, u))
if flags & DFSFlags.RETURN_DEPTHS:
depths[u] = depth
if (c := child[u]) < len(G[u]):
child[u] += 1
if (s := state[v := G[u][c]]) == 0: # Unvisited
if max_depth is None or depth <= max_depth:
if flags & DFSFlags.DOWN:
events.append((DFSEvent.DOWN, u, v))
stack[depth := depth+1] = v
if flags & DFSFlags.RETURN_PARENTS:
parents[v] = u
elif s == 1: # In progress
if flags & DFSFlags.BACK:
events.append((DFSEvent.BACK, u, v))
elif s == 2: # Completed
if flags & DFSFlags.CROSS:
events.append((DFSEvent.CROSS, u, v))
else:
depth -= 1
state[u] = 0 if DFSFlags.BACKTRACK in flags else 2
if flags & DFSFlags.LEAVE:
events.append((DFSEvent.LEAVE, u))
if depth != -1 and flags & DFSFlags.UP:
events.append((DFSEvent.UP, stack[depth], u))
if (DFSFlags.UP|DFSFlags.CONNECT_ROOTS) in flags:
events.append((DFSEvent.UP,-1,s))
ret = tuple((events,)) if DFSFlags.RETURN_ALL & flags else events
if DFSFlags.RETURN_PARENTS in flags:
ret += (parents,)
if DFSFlags.RETURN_DEPTHS in flags:
ret += (depths,)
return ret
def dfs_backtrack(G, flags: DFSFlags, s: Union[int,list] = None, max_depth: Union[int,None] = None):
stack_depth = (max_depth+1 if max_depth is not None else G.N)
stack = [0]*stack_depth
child = [0]*stack_depth
state = [0]*G.N
events: list[tuple[DFSEvent, int]|tuple[DFSEvent, int, int]] = []
for s in G.starts(s):
if state[s]: continue
state[s] = 1
stack[depth := 0] = s
if DFSFlags.DOWN|DFSFlags.CONNECT_ROOTS in flags:
events.append((DFSEvent.DOWN,-1,s))
while depth != -1:
u = stack[depth]
if state[u] == 1:
state[u] = 2
if DFSFlags.ENTER in flags:
events.append((DFSEvent.ENTER,u))
if max_depth is not None and depth >= max_depth:
child[depth] = len(G[u])
if DFSFlags.MAXDEPTH in flags:
events.append((DFSEvent.MAXDEPTH,u))
if (c := child[depth]) < len(G[u]):
child[depth] += 1
if state[v := G[u][c]]:
if DFSFlags.BACK in flags:
events.append((DFSEvent.BACK,u,v))
continue
state[v] = 1
if DFSFlags.DOWN in flags:
events.append((DFSEvent.DOWN,u,v))
stack[depth := depth+1] = v
else:
state[u] = 0
if DFSFlags.LEAVE in flags:
events.append((DFSEvent.LEAVE,u))
child[depth] = 0
depth -= 1
if depth and DFSFlags.UP in flags:
events.append((DFSEvent.UP, stack[depth], u))
if DFSFlags.UP|DFSFlags.CONNECT_ROOTS in flags:
events.append((DFSEvent.UP,-1,s))
return events
def dfs_enter_leave(G, s: Union[int,list,None] = None):
state = [True] * G.N
child: list[int] = elist(G.N)
stack: list[int] = elist(G.N)
events = []
for s in G.starts(s):
if not state[s]: continue
stack.append(s)
child.append(0)
while stack:
u = stack[-1]
if state[u]:
state[u] = False
events.append((DFSEvent.ENTER, u))
if (c := child[-1]) < len(G[u]):
child[-1] += 1
if state[v := G[u][c]]:
stack.append(v)
child.append(0)
else:
stack.pop()
child.pop()
events.append((DFSEvent.LEAVE, u))
return events
def dfs_topdown(G, s: Union[int,list,None] = None, connect_roots = False):
'''Returns list of (u,v) representing u->v edges in order of top down discovery'''
stack: list[int] = elist(G.N)
vis = [False]*G.N
edges: list[tuple[int,int]] = elist(G.N)
for s in G.starts(s):
if vis[s]: continue
if connect_roots:
edges.append((-1,s))
vis[s] = True
stack.append(s)
while stack:
u = stack.pop()
for v in G[u]:
if vis[v]: continue
vis[v] = True
edges.append((u,v))
stack.append(v)
return edges
def dfs_bottomup(G, s: Union[int,list,None] = None, connect_roots = False):
'''Returns list of (p,u) representing p->u edges in bottom up order'''
edges = G.dfs_topdown(s, connect_roots)
edges.reverse()
return edges
def is_bipartite(G):
N = G.N
que = deque()
color = [-1]*N
for s in range(N):
if color[s] >= 0:
continue
color[s] = 1
que.append(s)
while que:
u = que.popleft()
for v in G[u]:
if color[v] == -1:
color[v] = 1 - color[u]
que.append(v)
elif color[v] == color[u]:
return False
return True
def starts(G, v: Union[int,list,None]) -> Iterable:
if isinstance(v, int):
return (v,)
elif v is None:
return range(G.N)
else:
return v
@classmethod
def compile(cls, N: int, M: int, E):
edge = Parser.compile(E)
def parse(ts: TokenStream):
return cls(N, [edge(ts) for _ in range(M)])
return parse
class GraphWeightedProtocol(GraphProtocol):
def neighbors(G, v: int):
return map(operator.itemgetter(0), G[v])
@overload
def distance(G) -> list[list[int]]: ...
@overload
def distance(G, s: int = 0) -> list[int]: ...
@overload
def distance(G, s: int, g: int) -> int: ...
def distance(G, s = None, g = None):
if s == None:
return G.floyd_warshall()
else:
return G.dijkstra(s, g)
def dijkstra(G, s = 0, g = None):
D = [inf for _ in range(G.N)]
D[s] = 0
que = PriorityQueue(G.N)
que.push(s, 0)
while que:
v, d = que.pop()
if v == g: return d
if d > D[v]: continue
for c, w, *_ in G[v]:
if (nd := d + w) < D[c]:
D[c] = nd
que.push(c, nd)
return D if g is None else inf
@overload
def shortest_path(G, s: int, t: int) -> list[int]|None: ...
@overload
def shortest_path(G, s: int, t: int, distances = True) -> tuple[list[int]|None,list[int]]: ...
def shortest_path(G, s: int, t: int, distances = False):
D = [inf] * G.N
D[s] = 0
if s == t:
return ([], D) if distances else []
par = [-1] * G.N
down = [-1] * G.N
Eid = G.edge_ids()
que = PriorityQueue(G.N)
que.push(s, 0)
while que:
v, d = que.pop()
if v == t: break
if d > D[v]: continue
for i in range(len(G[v])):
c, w, *_ = G[v][i]
if (nd := d + w) < D[c]:
D[c] = nd
par[c] = v
down[c] = Eid[v][i]
que.push(c, nd)
if D[t] == inf:
return (None, D) if distances else None
path = []
v = t
while v != s:
path.append(down[v])
v = par[v]
return (path[::-1], D) if distances else path[::-1]
def kruskal(G):
E, N = G.E, G.N
heapify(E)
dsu = DSU(N)
MST = []
need = N-1
while E and need:
edge = heappop(E)
u,v,*_ = edge
u,v = dsu.merge(u,v,True)
if u != v:
MST.append(edge)
need -= 1
return MST
def bellman_ford(G, s = 0) -> list[int]:
D = [inf]*G.N
D[s] = 0
for _ in range(G.N-1):
for u, edges in enumerate(G):
if D[u] == inf: continue
for v,w,*_ in edges:
D[v] = min(D[v], D[u] + w)
return D
def floyd_warshall(G) -> list[list[int]]:
D = [[inf]*G.N for _ in range(G.N)]
for u, edges in enumerate(G):
D[u][u] = 0
for v,w in edges:
D[u][v] = min(D[u][v], w)
for k, Dk in enumerate(D):
for Di in D:
if Di[k] == inf: continue
for j in range(G.N):
if Dk[j] == inf: continue
Di[j] = min(Di[j], Di[k]+Dk[j])
return D
def dfs_events(G, flags: DFSFlags, s: Union[int,list,None] = None, max_depth: Union[int,None] = None):
if flags == DFSFlags.INTERVAL:
if max_depth is None:
return G.dfs_enter_leave(s)
elif flags == DFSFlags.DOWN or flags == DFSFlags.TOPDOWN:
if max_depth is None:
edges = G.dfs_topdown(s, DFSFlags.CONNECT_ROOTS in flags)
return [(DFSEvent.DOWN, p, u) for p,u in edges]
elif flags == DFSFlags.UP or flags == DFSFlags.BOTTOMUP:
if max_depth is None:
edges = G.dfs_bottomup(s, DFSFlags.CONNECT_ROOTS in flags)
return [(DFSEvent.UP, p, u) for p,u in edges]
elif flags & DFSFlags.BACKTRACK:
return G.dfs_backtrack(flags, s, max_depth)
state = [0] * G.N
child = elist(G.N)
weights = elist(G.N)
stack = elist(G.N)
if flags & DFSFlags.RETURN_PARENTS:
parents = [-1] * G.N
if flags & DFSFlags.RETURN_DEPTHS:
depths = [-1] * G.N
events = []
for s in G.starts(s):
stack.append(s)
child.append(0)
if (DFSFlags.DOWN|DFSFlags.CONNECT_ROOTS) in flags:
events.append((DFSEvent.DOWN,-1,s,-1))
while stack:
u = stack[-1]
if not state[u]:
state[u] = 1
if flags & DFSFlags.ENTER:
events.append((DFSEvent.ENTER, u))
if flags & DFSFlags.RETURN_DEPTHS:
depths[u] = len(stack)-1
if (c := child[-1]) < len(G[u]):
child[-1] += 1
v, w = G[u][c]
if (s := state[v]) == 0: # Unvisited
if max_depth is None or len(stack)-1 <= max_depth:
if flags & DFSFlags.DOWN:
events.append((DFSEvent.DOWN, u, v, w))
stack.append(v)
weights.append(w)
child.append(0)
if flags & DFSFlags.RETURN_PARENTS:
parents[v] = u
elif s == 1: # In progress
if flags & DFSFlags.BACK:
events.append((DFSEvent.BACK, u, v, w))
elif s == 2: # Completed
if flags & DFSFlags.CROSS:
events.append((DFSEvent.CROSS, u, v, w))
else:
stack.pop()
child.pop()
state[u] = 0 if DFSFlags.BACKTRACK in flags else 2
if flags & DFSFlags.LEAVE:
events.append((DFSEvent.LEAVE, u))
if stack and flags & DFSFlags.UP:
pw = weights.pop()
events.append((DFSEvent.UP, stack[-1], u, pw))
if (DFSFlags.UP|DFSFlags.CONNECT_ROOTS) in flags:
events.append((DFSEvent.UP,-1,s,-1))
ret = tuple((events,)) if DFSFlags.RETURN_ALL & flags else events
if DFSFlags.RETURN_PARENTS in flags:
ret += (parents,)
if DFSFlags.RETURN_DEPTHS in flags:
ret += (depths,)
return ret
def dfs_backtrack(G, flags: DFSFlags, s: Union[int,list] = None, max_depth: Union[int,None] = None):
stack_depth = (max_depth+1 if max_depth is not None else G.N)
stack = elist(stack_depth)
child = elist(stack_depth)
weights = elist(stack_depth)
state = [0]*G.N
events: list[tuple[DFSEvent, int]|tuple[DFSEvent, int, int]] = []
for s in G.starts(s):
if state[s]: continue
state[s] = 1
stack.append(s)
child.append(0)
if DFSFlags.DOWN|DFSFlags.CONNECT_ROOTS in flags:
events.append((DFSEvent.DOWN,-1,s,-1))
while stack:
u = stack[-1]
if state[u] == 1:
state[u] = 2
if DFSFlags.ENTER in flags:
events.append((DFSEvent.ENTER,u))
if max_depth is not None and len(stack) > max_depth:
child[-1] = len(G[u])
if DFSFlags.MAXDEPTH in flags:
events.append((DFSEvent.MAXDEPTH,u))
if (c := child[-1]) < len(G[u]):
child[-1] += 1
v, w = G[u][c]
if state[v]:
if DFSFlags.BACK in flags:
events.append((DFSEvent.BACK,u,v,w))
continue
state[v] = 1
if DFSFlags.DOWN in flags:
events.append((DFSEvent.DOWN,u,v,w))
stack.append(v)
child.append(0)
weights.append(w)
else:
state[u] = 0
if DFSFlags.LEAVE in flags:
events.append((DFSEvent.LEAVE,u))
stack.pop()
child.pop()
if stack and DFSFlags.UP in flags:
pw = weights.pop()
events.append((DFSEvent.UP, stack[-1], u, pw))
if DFSFlags.UP|DFSFlags.CONNECT_ROOTS in flags:
events.append((DFSEvent.UP,-1,s,-1))
return events
def dfs_topdown(G, s: Union[int,list[int],None] = None, connect_roots = False):
'''Returns list of (u,v) representing u->v edges in order of top down discovery'''
stack: list[int] = elist(G.N)
vis = [False]*G.N
edges: list[tuple[int,int]] = elist(G.N)
for s in G.starts(s):
if vis[s]: continue
if connect_roots:
edges.append((-1,s,-1))
vis[s] = True
stack.append(s)
while stack:
u = stack.pop()
for v,w in G[u]:
if vis[v]: continue
vis[v] = True
edges.append((u,v,w))
stack.append(v)
return edges
from typing import Sequence
class CSRIncremental(Sequence[list[_T]]):
def __init__(csr, sizes: list[int]):
csr.L, N = [0]*len(sizes), 0
for i,sz in enumerate(sizes):
csr.L[i] = N; N += sz
csr.R, csr.A = csr.L[:], [0]*N
def append(csr, i: int, x: _T):
csr.A[csr.R[i]] = x; csr.R[i] += 1
def __iter__(csr):
for i,l in enumerate(csr.L):
yield csr.A[l:csr.R[i]]
def __getitem__(csr, i: int) -> _T:
return csr.A[i]
def __len__(dsu):
return len(dsu.L)
def range(csr, i: int) -> _T:
return range(csr.L[i], csr.R[i])
class DSU(Parsable, Collection):
def __init__(dsu, N):
dsu.N, dsu.cc, dsu.par = N, N, [-1]*N
def merge(dsu, u, v, src = False):
x, y = dsu.leader(u), dsu.leader(v)
if x == y: return (x,y) if src else x
if dsu.par[x] > dsu.par[y]: x, y = y, x
dsu.par[x] += dsu.par[y]; dsu.par[y] = x; dsu.cc -= 1
return (x,y) if src else x
def same(dsu, u: int, v: int):
return dsu.leader(u) == dsu.leader(v)
def leader(dsu, i) -> int:
p = (par := dsu.par)[i]
while p >= 0:
if par[p] < 0: return p
par[i], i, p = par[p], par[p], par[par[p]]
return i
def size(dsu, i) -> int:
return -dsu.par[dsu.leader(i)]
def groups(dsu) -> CSRIncremental[int]:
sizes, row, p = [0]*dsu.cc, [-1]*dsu.N, 0
for i in range(dsu.cc):
while dsu.par[p] >= 0: p += 1
sizes[i], row[p] = -dsu.par[p], i; p += 1
csr = CSRIncremental(sizes)
for i in range(dsu.N): csr.append(row[dsu.leader(i)], i)
return csr
__iter__ = groups
def __len__(dsu):
return dsu.cc
def __contains__(dsu, uv):
u, v = uv
return dsu.same(u, v)
@classmethod
def compile(cls, N: int, M: int, shift = -1):
def parse_fn(ts: TokenStream):
dsu = cls(N)
for _ in range(M):
u, v = ts._line()
dsu.merge(int(u)+shift, int(v)+shift)
return dsu
return parse_fn
from collections import UserList
from heapq import heapify, heappop, heappush, heappushpop, heapreplace
from typing import Generic
class HeapProtocol(Generic[_T]):
def pop(self) -> _T: ...
def push(self, item: _T): ...
def pushpop(self, item: _T) -> _T: ...
def replace(self, item: _T) -> _T: ...
class PriorityQueue(HeapProtocol[int], UserList[int]):
def __init__(self, N: int, ids: list[int] = None, priorities: list[int] = None, /):
self.shift = N.bit_length()
self.mask = (1 << self.shift)-1
if ids is None:
self.data = elist(N)
elif priorities is None:
heapify(ids)
self.data = ids
else:
M = len(ids)
data = [0]*M
for i in range(M):
data[i] = self.encode(ids[i], priorities[i])
heapify(data)
self.data = data
def encode(self, id, priority):
return priority << self.shift | id
def decode(self, encoded):
return self.mask & encoded, encoded >> self.shift
def pop(self):
return self.decode(heappop(self.data))
def push(self, id: int, priority: int):
heappush(self.data, self.encode(id, priority))
def pushpop(self, id: int, priority: int):
return self.decode(heappushpop(self.data, self.encode(id, priority)))
def replace(self, id: int, priority: int):
return self.decode(heapreplace(self.data, self.encode(id, priority)))
from operator import itemgetter
class DiGraphWeighted(GraphWeightedProtocol):
def __init__(G, N, E: list = []):
super().__init__(N, E, ([] for _ in range(N)))
for u,v,*w in G.E:
G[u].append((v,*w))
def edge_ids(G) -> list[list[int]]:
Eid = [[] for _ in range(G.N)]
for e,(u,v,*w) in enumerate(G.E):
Eid[u].append(e)
return Eid
def neighbors(G, v: int) -> Iterable[int]:
return map(itemgetter(0), G[v])
@classmethod
def compile(cls, N: int, M: int, E: Union[type,int] = EdgeWeighted[-1]):
if isinstance(E, int): E = EdgeWeighted[E]
return super().compile(N, M, E)
class DiGraph(GraphProtocol):
def __init__(G, N: int, E: list[Edge]=[]):
super().__init__(N, E, ([] for _ in range(N)))
for u,v in G.E:
G[u].append(v)
def edge_ids(G) -> list[list[int]]:
Eid = [[] for _ in range(G.N)]
for e,(u,v) in enumerate(G.E):
Eid[u].append(e)
return Eid
@classmethod
def compile(cls, N: int, M: int, E: Union[type,int] = Edge[-1]):
if isinstance(E, int): E = Edge[E]
return super().compile(N, M, E)
from typing import Iterable, Type, Union, overload
@overload
def read() -> Iterable[int]: ...
@overload
def read(spec: int) -> list[int]: ...
@overload
def read(spec: Union[Type[_T],_T], char=False) -> _T: ...
def read(spec: Union[Type[_T],_T] = None, char=False):
if not char and spec is None: return map(int, TokenStream.default.line())
parser: _T = Parser.compile(spec)
return parser(CharStream.default if char else TokenStream.default)
def write(*args, **kwargs):
'''Prints the values to a stream, or to stdout_fast by default.'''
sep, file = kwargs.pop("sep", " "), kwargs.pop("file", IOWrapper.stdout)
at_start = True
for x in args:
if not at_start:
file.write(sep)
file.write(str(x))
at_start = False
file.write(kwargs.pop("end", "\n"))
if kwargs.pop("flush", False):
file.flush()
if __name__ == "__main__":
main()