cp-library
This documentation is automatically generated by online-judge-tools/verification-helper
test/library-checker/graph/directedmst_skew_heap_forest.test.py
- View this file on GitHub
- Last update: 2025-06-20 03:24:59+09:00
- Problem: https://judge.yosupo.jp/problem/directedmst
Depends on
Code
# verification-helper: PROBLEM https://judge.yosupo.jp/problem/directedmst
import os,sys,io
input=io.BytesIO(os.read(0,os.fstat(0).st_size)).readline
class UnionFind:
def __init__(self, n):
self.parent = [-1] * n
self.n = n
def root(self, x):
if self.parent[x] < 0:
return x
else:
self.parent[x] = self.root(self.parent[x])
return self.parent[x]
def merge(self, x, y):
x = self.root(x)
y = self.root(y)
if x == y:
return False
self.parent[x] += self.parent[y]
self.parent[y] = x
return True
def same(self, x, y):
return self.root(x) == self.root(y)
from cp_library.ds.heap.skew_heap_forest_cls import SkewHeapForest
def directed_mst(n, edges, root):
OFFSET = 1 << 31
from_cost = [0] * n
from_heap = SkewHeapForest(n, len(edges))
from_ = [0] * n
uf = UnionFind(n)
par_e = [-1] * m
stem = [-1] * n
used = [0] * n
used[root] = 2
idxs = []
for idx, (fr, to, cost) in enumerate(edges):
from_heap.push(to, cost << 31 | idx)
res = 0
for v in range(n):
if used[v] != 0:
continue
processing = []
chi_e = []
cycle = 0
while used[v] != 2:
used[v] = 1
processing.append(v)
if from_heap.empty(v):
return -1, par
from_cost[v], idx = divmod(from_heap.pop(v), OFFSET)
from_[v] = uf.root(edges[idx][0])
if stem[v] == -1:
stem[v] = idx
if from_[v] == v:
continue
res += from_cost[v]
idxs.append(idx)
while cycle:
par_e[chi_e.pop()] = idx
cycle -= 1
chi_e.append(idx)
if used[from_[v]] == 1:
p = v
while True:
if not from_heap.empty(p):
from_heap.add(p, -from_cost[p] << 31)
if p != v:
uf.merge(v, p)
from_heap.roots[v] = from_heap.merge(from_heap.roots[v], from_heap.roots[p])
p = uf.root(from_[p])
cycle += 1
if p == v:
break
else:
v = from_[v]
for v in processing:
used[v] = 2
used_e = [0] * m
tree = [-1] * n
for idx in reversed(idxs):
if used_e[idx]:
continue
fr, to, cost = edges[idx]
tree[to] = fr
x = stem[to]
while x != idx:
used_e[x] = 1
x = par_e[x]
return res, tree
n, m, root = map(int, input().split())
edges = [[int(s) for s in input().split()] for i in range(m)]
res, par = directed_mst(n, edges, root)
if res == -1:
print(res)
else:
print(res)
print(*[p if p != -1 else i for i, p in enumerate(par)])# verification-helper: PROBLEM https://judge.yosupo.jp/problem/directedmst
import os,sys,io
input=io.BytesIO(os.read(0,os.fstat(0).st_size)).readline
class UnionFind:
def __init__(self, n):
self.parent = [-1] * n
self.n = n
def root(self, x):
if self.parent[x] < 0:
return x
else:
self.parent[x] = self.root(self.parent[x])
return self.parent[x]
def merge(self, x, y):
x = self.root(x)
y = self.root(y)
if x == y:
return False
self.parent[x] += self.parent[y]
self.parent[y] = x
return True
def same(self, x, y):
return self.root(x) == self.root(y)
'''
╺━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╸
https://kobejean.github.io/cp-library
'''
import operator
from typing import Generic
from typing import TypeVar
_T = TypeVar('T')
_U = TypeVar('U')
def elist(est_len: int) -> list: ...
try:
from __pypy__ import newlist_hint
except:
def newlist_hint(hint):
return []
elist = newlist_hint
class SkewHeapForest(Generic[_T]):
def __init__(shf, N, M, e: _T = 0, op = operator.add):
shf.V, shf.A, shf.L, shf.R, shf.roots = [e]*M, [e]*M, [-1]*M, [-1]*M, [-1]*N
shf.id, shf.st, shf.e, shf.op = 0, elist(M), e, op
def propagate(shf, u: int):
if (a := shf.A[u]) != shf.e:
if ~(l := shf.L[u]): shf.A[l] = shf.op(shf.A[l], a)
if ~(r := shf.R[u]): shf.A[r] = shf.op(shf.A[r], a)
shf.V[u] = shf.op(shf.V[u], a); shf.A[u] = shf.e
def merge(shf, u: int, v: int):
while ~u and ~v:
shf.propagate(u); shf.propagate(v)
if shf.V[v] < shf.V[u]: u, v = v, u
shf.st.append(u); shf.R[u], u = shf.L[u], shf.R[u]
u = u if ~u else v
while shf.st: shf.L[u := shf.st.pop()] = u
return u
def min(shf, i: int):
assert ~(root := shf.roots[i])
shf.propagate(root)
return shf.V[root]
def push(shf, i: int, x: _T):
shf.id = (id := shf.id)+1
shf.V[id] = x
shf.roots[i] = shf.merge(shf.roots[i], id)
def pop(shf, i: int) -> _T:
assert ~(root := shf.roots[i])
shf.propagate(root)
val, shf.roots[i] = shf.V[root], shf.merge(shf.L[root], shf.R[root])
return val
def add(shf, i: int, val: _T): shf.A[shf.roots[i]] = shf.op(shf.A[shf.roots[i]], val)
def empty(shf, i: int): return shf.roots[i] == -1
def directed_mst(n, edges, root):
OFFSET = 1 << 31
from_cost = [0] * n
from_heap = SkewHeapForest(n, len(edges))
from_ = [0] * n
uf = UnionFind(n)
par_e = [-1] * m
stem = [-1] * n
used = [0] * n
used[root] = 2
idxs = []
for idx, (fr, to, cost) in enumerate(edges):
from_heap.push(to, cost << 31 | idx)
res = 0
for v in range(n):
if used[v] != 0:
continue
processing = []
chi_e = []
cycle = 0
while used[v] != 2:
used[v] = 1
processing.append(v)
if from_heap.empty(v):
return -1, par
from_cost[v], idx = divmod(from_heap.pop(v), OFFSET)
from_[v] = uf.root(edges[idx][0])
if stem[v] == -1:
stem[v] = idx
if from_[v] == v:
continue
res += from_cost[v]
idxs.append(idx)
while cycle:
par_e[chi_e.pop()] = idx
cycle -= 1
chi_e.append(idx)
if used[from_[v]] == 1:
p = v
while True:
if not from_heap.empty(p):
from_heap.add(p, -from_cost[p] << 31)
if p != v:
uf.merge(v, p)
from_heap.roots[v] = from_heap.merge(from_heap.roots[v], from_heap.roots[p])
p = uf.root(from_[p])
cycle += 1
if p == v:
break
else:
v = from_[v]
for v in processing:
used[v] = 2
used_e = [0] * m
tree = [-1] * n
for idx in reversed(idxs):
if used_e[idx]:
continue
fr, to, cost = edges[idx]
tree[to] = fr
x = stem[to]
while x != idx:
used_e[x] = 1
x = par_e[x]
return res, tree
n, m, root = map(int, input().split())
edges = [[int(s) for s in input().split()] for i in range(m)]
res, par = directed_mst(n, edges, root)
if res == -1:
print(res)
else:
print(res)
print(*[p if p != -1 else i for i, p in enumerate(par)])