#include "graph/minimum-cost-arborescence.hpp"
#pragma once #include "../data-structure/skew-heap.hpp" #include "../data-structure/union-find.hpp" #include "graph-template.hpp" template <typename T> Edges<T> MinimumCostArborescence(int N, int root, const Edges<T> &es) { using Heap = SkewHeap<T>; using Ptr = typename Heap::Ptr; UnionFind uf(N); vector<int> used(N, -1), from(N); vector<T> from_cost(N); vector<Ptr> come(N, nullptr); used[root] = root; vector<int> par_e(es.size(), -1), stem(N, -1), idxs; for (int i = 0; i < (int)es.size(); i++) { auto &e = es[i]; come[e] = Heap::push(come[e], e.cost, i); } T costs = 0; for (int start = 0; start < N; start++) { if (used[start] != -1) continue; int cur = start; vector<int> chi_e; int cycle = 0; while (used[cur] == -1 || used[cur] == start) { used[cur] = start; if (come[cur] == nullptr) return {}; int src = uf.find(es[come[cur]->idx].src); T cost = come[cur]->key + come[cur]->laz; int idx = come[cur]->idx; come[cur] = Heap::pop(come[cur]); if (src == cur) continue; from[cur] = src; from_cost[cur] = cost; if (stem[cur] == -1) stem[cur] = idx; costs += cost; idxs.push_back(idx); while (cycle) { par_e[chi_e.back()] = idx; chi_e.pop_back(); --cycle; } chi_e.push_back(idx); if (used[src] == start) { int p = cur; do { if (come[p]) Heap::apply(come[p], -from_cost[p]); if (p != cur) { uf.unite(p, cur); Ptr newheap = Heap::meld(come[cur], come[p]); come[cur = uf.find(cur)] = newheap; } p = uf.find(from[p]); ++cycle; } while (p != cur); } else { cur = src; } } } vector<int> used_e(es.size()); Edges<T> res; for (int _ = (int)idxs.size(); _--;) { int idx = idxs[_]; if (used_e[idx]) continue; auto &e = es[idx]; res.push_back(e); int x = stem[e.to]; while (x != idx) { used_e[x] = true; x = par_e[x]; } } return res; }
#line 2 "graph/minimum-cost-arborescence.hpp" #line 2 "data-structure/skew-heap.hpp" template <typename T, bool isMin = true> struct SkewHeap { struct Node { T key, laz; Node *l, *r; int idx; Node() = default; Node(const T &k, int i = -1) : key(k), laz(0), l(nullptr), r(nullptr), idx(i) {} }; using Ptr = Node *; static void propagate(Ptr x) { if (x->laz == 0) return; if (x->l) x->l->laz += x->laz; if (x->r) x->r->laz += x->laz; x->key += x->laz; x->laz = 0; } static Ptr meld(Ptr x, Ptr y) { if (!x || !y) return x ? x : y; if (!comp(x, y)) swap(x, y); propagate(x); x->r = meld(x->r, y); swap(x->l, x->r); return x; } static Ptr alloc(const T &key, int idx = -1) { return new Node(key, idx); } static Ptr pop(Ptr x) { propagate(x); return meld(x->l, x->r); } static Ptr push(Ptr x, const T &key, int idx = -1) { return meld(x, alloc(key, idx)); } static void apply(Ptr x, const T &laz) { x->laz += laz; propagate(x); } private: static inline bool comp(Ptr x, Ptr y) { if constexpr (isMin) { return x->key + x->laz < y->key + y->laz; } else { return x->key + x->laz > y->key + y->laz; } } }; /** * @brief Skew Heap * @docs docs/data-structure/skew-heap.md */ #line 2 "data-structure/union-find.hpp" struct UnionFind { vector<int> data; UnionFind(int N) : data(N, -1) {} int find(int k) { return data[k] < 0 ? k : data[k] = find(data[k]); } int unite(int x, int y) { if ((x = find(x)) == (y = find(y))) return false; if (data[x] > data[y]) swap(x, y); data[x] += data[y]; data[y] = x; return true; } // f ... merge function template<typename F> int unite(int x, int y,const F &f) { if ((x = find(x)) == (y = find(y))) return false; if (data[x] > data[y]) swap(x, y); data[x] += data[y]; data[y] = x; f(x, y); return true; } int size(int k) { return -data[find(k)]; } int same(int x, int y) { return find(x) == find(y); } }; /** * @brief Union Find(Disjoint Set Union) * @docs docs/data-structure/union-find.md */ #line 2 "graph/graph-template.hpp" template <typename T> struct edge { int src, to; T cost; edge(int _to, T _cost) : src(-1), to(_to), cost(_cost) {} edge(int _src, int _to, T _cost) : src(_src), to(_to), cost(_cost) {} edge &operator=(const int &x) { to = x; return *this; } operator int() const { return to; } }; template <typename T> using Edges = vector<edge<T>>; template <typename T> using WeightedGraph = vector<Edges<T>>; using UnweightedGraph = vector<vector<int>>; // Input of (Unweighted) Graph UnweightedGraph graph(int N, int M = -1, bool is_directed = false, bool is_1origin = true) { UnweightedGraph g(N); if (M == -1) M = N - 1; for (int _ = 0; _ < M; _++) { int x, y; cin >> x >> y; if (is_1origin) x--, y--; g[x].push_back(y); if (!is_directed) g[y].push_back(x); } return g; } // Input of Weighted Graph template <typename T> WeightedGraph<T> wgraph(int N, int M = -1, bool is_directed = false, bool is_1origin = true) { WeightedGraph<T> g(N); if (M == -1) M = N - 1; for (int _ = 0; _ < M; _++) { int x, y; cin >> x >> y; T c; cin >> c; if (is_1origin) x--, y--; g[x].emplace_back(x, y, c); if (!is_directed) g[y].emplace_back(y, x, c); } return g; } // Input of Edges template <typename T> Edges<T> esgraph(int N, int M, int is_weighted = true, bool is_1origin = true) { Edges<T> es; for (int _ = 0; _ < M; _++) { int x, y; cin >> x >> y; T c; if (is_weighted) cin >> c; else c = 1; if (is_1origin) x--, y--; es.emplace_back(x, y, c); } return es; } // Input of Adjacency Matrix template <typename T> vector<vector<T>> adjgraph(int N, int M, T INF, int is_weighted = true, bool is_directed = false, bool is_1origin = true) { vector<vector<T>> d(N, vector<T>(N, INF)); for (int _ = 0; _ < M; _++) { int x, y; cin >> x >> y; T c; if (is_weighted) cin >> c; else c = 1; if (is_1origin) x--, y--; d[x][y] = c; if (!is_directed) d[y][x] = c; } return d; } /** * @brief グラフテンプレート * @docs docs/graph/graph-template.md */ #line 6 "graph/minimum-cost-arborescence.hpp" template <typename T> Edges<T> MinimumCostArborescence(int N, int root, const Edges<T> &es) { using Heap = SkewHeap<T>; using Ptr = typename Heap::Ptr; UnionFind uf(N); vector<int> used(N, -1), from(N); vector<T> from_cost(N); vector<Ptr> come(N, nullptr); used[root] = root; vector<int> par_e(es.size(), -1), stem(N, -1), idxs; for (int i = 0; i < (int)es.size(); i++) { auto &e = es[i]; come[e] = Heap::push(come[e], e.cost, i); } T costs = 0; for (int start = 0; start < N; start++) { if (used[start] != -1) continue; int cur = start; vector<int> chi_e; int cycle = 0; while (used[cur] == -1 || used[cur] == start) { used[cur] = start; if (come[cur] == nullptr) return {}; int src = uf.find(es[come[cur]->idx].src); T cost = come[cur]->key + come[cur]->laz; int idx = come[cur]->idx; come[cur] = Heap::pop(come[cur]); if (src == cur) continue; from[cur] = src; from_cost[cur] = cost; if (stem[cur] == -1) stem[cur] = idx; costs += cost; idxs.push_back(idx); while (cycle) { par_e[chi_e.back()] = idx; chi_e.pop_back(); --cycle; } chi_e.push_back(idx); if (used[src] == start) { int p = cur; do { if (come[p]) Heap::apply(come[p], -from_cost[p]); if (p != cur) { uf.unite(p, cur); Ptr newheap = Heap::meld(come[cur], come[p]); come[cur = uf.find(cur)] = newheap; } p = uf.find(from[p]); ++cycle; } while (p != cur); } else { cur = src; } } } vector<int> used_e(es.size()); Edges<T> res; for (int _ = (int)idxs.size(); _--;) { int idx = idxs[_]; if (used_e[idx]) continue; auto &e = es[idx]; res.push_back(e); int x = stem[e.to]; while (x != idx) { used_e[x] = true; x = par_e[x]; } } return res; }