This documentation is automatically generated by online-judge-tools/verification-helper
View the Project on GitHub ebi-fly13/Library
#include "graph/directed_mst.hpp"
根と有向グラフを与えて、有向最小全域木を構築する。返り値 $(x, p)$ について、 $x$ は木の重みの総和、 $p_{i}$ は頂点 $i$ の親である。 $p_{root} = root$ である。 $O(M\log {N})$
#pragma once #include <cassert> #include <numeric> #include <optional> #include <ranges> #include <utility> #include <vector> #include "../data_structure/skew_heap.hpp" #include "../graph/base.hpp" namespace ebi { namespace internal { struct union_find_not_merge_tech { public: union_find_not_merge_tech(int n_) : n(n_), data(n, -1) {} int leader(int u) { if (data[u] < 0) return u; return data[u] = leader(data[u]); } bool merge(int u, int v) { u = leader(u); v = leader(v); if (u == v) return false; data[u] += data[v]; data[v] = u; return true; } bool same(int u, int v) { return leader(u) == leader(v); } int size(int u) { u = leader(u); return -data[u]; } private: int n; std::vector<int> data; }; } // namespace internal template <class T> std::optional<std::pair<T, std::vector<int>>> directed_mst(int root, const Graph<T> &g) { using Heap = skew_heap<T, Edge<T>, std::greater<T>>; int n = g.node_number(); std::vector<int> used(2 * n, 0); std::vector<int> par_node(2 * n, -1); std::vector<int> best_id(2 * n, -1); std::vector<T> best_cost(2 * n); internal::union_find_not_merge_tech uf(2 * n); std::vector<Heap> heap(2 * n); int nxt = n; used[root] = 2; for (auto e : g.get_edges()) { heap[e.to].push(e.cost, e); } T sum = 0; for (auto s : std::views::iota(0, n)) { if (used[s]) continue; std::vector<int> path = {s}; while (1) { int a = path.back(); if (heap[a].empty()) return std::nullopt; auto [c, e] = heap[a].top(); heap[a].pop(); int v = uf.leader(e.from); if (a == v) continue; used[a] = 1; best_id[a] = e.id; best_cost[a] = c; sum += c; if (!used[v]) { path.emplace_back(v); continue; } else if (used[v] == 1) { int w = -1; int u = nxt++; while (w != v) { w = path.back(); path.pop_back(); T sub = best_cost[w]; heap[w].add(-sub); heap[u].meld(heap[w]); par_node[w] = u; uf.merge(u, w); used[w] = 2; } path.emplace_back(u); } else break; } for (auto v : path) used[v] = 2; } std::vector<int> par(n, -1); std::vector<bool> done(nxt, false); done[root] = true; par[root] = root; T ret = 0; for (auto i : std::views::iota(0, nxt) | std::views::reverse) { if (done[i]) continue; int id = best_id[i]; auto e = g.get_edge(id); par[e.to] = e.from; ret += e.cost; int x = e.to; while (x != -1 && !done[x]) { done[x] = true; x = par_node[x]; } } assert(sum == ret); return std::pair<T, std::vector<int>>{sum, par}; } } // namespace ebi
#line 2 "graph/directed_mst.hpp" #include <cassert> #include <numeric> #include <optional> #include <ranges> #include <utility> #include <vector> #line 2 "data_structure/skew_heap.hpp" #line 4 "data_structure/skew_heap.hpp" #include <functional> #include <memory> #line 7 "data_structure/skew_heap.hpp" namespace ebi { template <class Key, class T, class Compare = std::less<Key>> struct skew_heap { private: using value_type = Key; using Self = skew_heap<Key, T, Compare>; struct Node; using iterator = std::shared_ptr<Node>; struct Node { Node(value_type x_, T info_) : x(x_), info(info_) {} void propagate() { if (lhs) lhs->lazy += lazy; if (rhs) rhs->lazy += lazy; x += lazy; lazy = value_type(); } value_type x; T info; value_type lazy = value_type(); iterator lhs = nullptr, rhs = nullptr; }; iterator internal_meld(iterator lhs, iterator rhs) { if (lhs == nullptr) return rhs; if (rhs == nullptr) return lhs; lhs->propagate(); rhs->propagate(); if (Compare()(lhs->x, rhs->x)) { std::swap(lhs, rhs); } lhs->rhs = internal_meld(lhs->rhs, rhs); std::swap(lhs->lhs, lhs->rhs); return lhs; } public: skew_heap() = default; void push(value_type x, T info) { root = internal_meld(root, std::make_shared<Node>(x, info)); sz++; } void pop() { assert(!empty()); root = internal_meld(root->lhs, root->rhs); sz--; } void meld(Self &heap) { root = internal_meld(root, heap.root); sz += heap.sz; } void add(value_type lazy) { if (root == nullptr) return; root->lazy += lazy; root->propagate(); } bool empty() const { return root == nullptr; } std::pair<value_type, T> top() const { return {root->x, root->info}; } int size() const { return sz; } private: iterator root = nullptr; int sz = 0; }; } // namespace ebi #line 2 "graph/base.hpp" #line 4 "graph/base.hpp" #include <iostream> #line 7 "graph/base.hpp" #line 2 "data_structure/simple_csr.hpp" #line 6 "data_structure/simple_csr.hpp" namespace ebi { template <class E> struct simple_csr { simple_csr() = default; simple_csr(int n, const std::vector<std::pair<int, E>>& elements) : start(n + 1, 0), elist(elements.size()) { for (auto e : elements) { start[e.first + 1]++; } for (auto i : std::views::iota(0, n)) { start[i + 1] += start[i]; } auto counter = start; for (auto [i, e] : elements) { elist[counter[i]++] = e; } } simple_csr(const std::vector<std::vector<E>>& es) : start(es.size() + 1, 0) { int n = es.size(); for (auto i : std::views::iota(0, n)) { start[i + 1] = (int)es[i].size() + start[i]; } elist.resize(start.back()); for (auto i : std::views::iota(0, n)) { std::copy(es[i].begin(), es[i].end(), elist.begin() + start[i]); } } int size() const { return (int)start.size() - 1; } const auto operator[](int i) const { return std::ranges::subrange(elist.begin() + start[i], elist.begin() + start[i + 1]); } auto operator[](int i) { return std::ranges::subrange(elist.begin() + start[i], elist.begin() + start[i + 1]); } const auto operator()(int i, int l, int r) const { return std::ranges::subrange(elist.begin() + start[i] + l, elist.begin() + start[i + 1] + r); } auto operator()(int i, int l, int r) { return std::ranges::subrange(elist.begin() + start[i] + l, elist.begin() + start[i + 1] + r); } private: std::vector<int> start; std::vector<E> elist; }; } // namespace ebi #line 9 "graph/base.hpp" namespace ebi { template <class T> struct Edge { int from, to; T cost; int id; }; template <class E> struct Graph { using cost_type = E; using edge_type = Edge<cost_type>; Graph(int n_) : n(n_) {} Graph() = default; void add_edge(int u, int v, cost_type c) { buff.emplace_back(u, edge_type{u, v, c, m}); edges.emplace_back(edge_type{u, v, c, m++}); } void add_undirected_edge(int u, int v, cost_type c) { buff.emplace_back(u, edge_type{u, v, c, m}); buff.emplace_back(v, edge_type{v, u, c, m}); edges.emplace_back(edge_type{u, v, c, m}); m++; } void read_tree(int offset = 1, bool is_weighted = false) { read_graph(n - 1, offset, false, is_weighted); } void read_parents(int offset = 1) { for (auto i : std::views::iota(1, n)) { int p; std::cin >> p; p -= offset; add_undirected_edge(p, i, 1); } build(); } void read_graph(int e, int offset = 1, bool is_directed = false, bool is_weighted = false) { for (int i = 0; i < e; i++) { int u, v; std::cin >> u >> v; u -= offset; v -= offset; if (is_weighted) { cost_type c; std::cin >> c; if (is_directed) { add_edge(u, v, c); } else { add_undirected_edge(u, v, c); } } else { if (is_directed) { add_edge(u, v, 1); } else { add_undirected_edge(u, v, 1); } } } build(); } void build() { assert(!prepared); csr = simple_csr<edge_type>(n, buff); buff.clear(); prepared = true; } int size() const { return n; } int node_number() const { return n; } int edge_number() const { return m; } edge_type get_edge(int i) const { return edges[i]; } std::vector<edge_type> get_edges() const { return edges; } const auto operator[](int i) const { return csr[i]; } auto operator[](int i) { return csr[i]; } private: int n, m = 0; std::vector<std::pair<int,edge_type>> buff; std::vector<edge_type> edges; simple_csr<edge_type> csr; bool prepared = false; }; } // namespace ebi #line 12 "graph/directed_mst.hpp" namespace ebi { namespace internal { struct union_find_not_merge_tech { public: union_find_not_merge_tech(int n_) : n(n_), data(n, -1) {} int leader(int u) { if (data[u] < 0) return u; return data[u] = leader(data[u]); } bool merge(int u, int v) { u = leader(u); v = leader(v); if (u == v) return false; data[u] += data[v]; data[v] = u; return true; } bool same(int u, int v) { return leader(u) == leader(v); } int size(int u) { u = leader(u); return -data[u]; } private: int n; std::vector<int> data; }; } // namespace internal template <class T> std::optional<std::pair<T, std::vector<int>>> directed_mst(int root, const Graph<T> &g) { using Heap = skew_heap<T, Edge<T>, std::greater<T>>; int n = g.node_number(); std::vector<int> used(2 * n, 0); std::vector<int> par_node(2 * n, -1); std::vector<int> best_id(2 * n, -1); std::vector<T> best_cost(2 * n); internal::union_find_not_merge_tech uf(2 * n); std::vector<Heap> heap(2 * n); int nxt = n; used[root] = 2; for (auto e : g.get_edges()) { heap[e.to].push(e.cost, e); } T sum = 0; for (auto s : std::views::iota(0, n)) { if (used[s]) continue; std::vector<int> path = {s}; while (1) { int a = path.back(); if (heap[a].empty()) return std::nullopt; auto [c, e] = heap[a].top(); heap[a].pop(); int v = uf.leader(e.from); if (a == v) continue; used[a] = 1; best_id[a] = e.id; best_cost[a] = c; sum += c; if (!used[v]) { path.emplace_back(v); continue; } else if (used[v] == 1) { int w = -1; int u = nxt++; while (w != v) { w = path.back(); path.pop_back(); T sub = best_cost[w]; heap[w].add(-sub); heap[u].meld(heap[w]); par_node[w] = u; uf.merge(u, w); used[w] = 2; } path.emplace_back(u); } else break; } for (auto v : path) used[v] = 2; } std::vector<int> par(n, -1); std::vector<bool> done(nxt, false); done[root] = true; par[root] = root; T ret = 0; for (auto i : std::views::iota(0, nxt) | std::views::reverse) { if (done[i]) continue; int id = best_id[i]; auto e = g.get_edge(id); par[e.to] = e.from; ret += e.cost; int x = e.to; while (x != -1 && !done[x]) { done[x] = true; x = par_node[x]; } } assert(sum == ret); return std::pair<T, std::vector<int>>{sum, par}; } } // namespace ebi