Merge pull request #385 from hitonanode/01-on-tree

01 on Tree

Author: hitonanode

other_algorithms/test/tree_pop_order_optimization.test.cpp

@@ -0,0 +1,30 @@
+#define PROBLEM                                                                                   \
+    "https://judge.yosupo.jp/problem/rooted_tree_topological_order_with_minimum_inversions"
+#include "../tree_pop_order_optimization.hpp"
+#include <iostream>
+#include <vector>
+using namespace std;
+
+int main() {
+    cin.tie(nullptr), ios::sync_with_stdio(false);
+
+    int N;
+    cin >> N;
+    vector<vector<int>> to(N);
+
+    for (int i = 1; i < N; ++i) {
+        int p;
+        cin >> p;
+        to.at(p).push_back(i);
+        to.at(i).push_back(p);
+    }
+
+    vector<long long> c(N), d(N);
+    for (auto &e : c) cin >> e;
+    for (auto &e : d) cin >> e;
+
+    const auto [order, ret] = Solve01OnTree(to, c, d, 0);
+    cout << ret << '\n';
+    for (auto e : order) cout << e << ' ';
+    cout << '\n';
+}

other_algorithms/test/tree_pop_order_optimization.yuki3148.test.cpp

@@ -0,0 +1,48 @@
+#define PROBLEM "https://yukicoder.me/problems/no/3148"
+#include "../tree_pop_order_optimization.hpp"
+
+#include <algorithm>
+#include <iostream>
+#include <vector>
+using namespace std;
+
+int main() {
+    cin.tie(nullptr), ios::sync_with_stdio(false);
+
+    int N;
+    string str;
+    cin >> N >> str;
+    vector<vector<int>> child(N + 1);
+    {
+        vector<int> stk{0};
+        int openid = 1;
+        for (auto c : str) {
+            if (c == '(') {
+                stk.push_back(openid++);
+            } else {
+                int v = stk.back();
+                stk.pop_back();
+                if (stk.size()) child.at(stk.back()).push_back(v);
+            }
+        }
+    }
+
+    vector<long long> A(N);
+    for (auto &a : A) cin >> a;
+    A.insert(A.begin(), 0);
+
+    vector<long long> B(A.size(), 1);
+    B.at(0) = 0;
+
+    vector<int> seq = Solve01OnTree(child, A, B, 0).first;
+    std::reverse(seq.begin(), seq.end());
+
+    long long dy = 0, ans = 0;
+    for (int i : seq) {
+        if (i == 0) continue;
+        dy += A.at(i);
+        ans += dy;
+    }
+
+    cout << ans << '\n';
+}

other_algorithms/tree_pop_order_optimization.hpp

@@ -0,0 +1,128 @@
+#pragma once
+#include <algorithm>
+#include <cassert>
+#include <queue>
+#include <utility>
+#include <vector>
+
+// "01 on Tree"
+// https://judge.yosupo.jp/problem/rooted_tree_topological_order_with_minimum_inversions
+// https://yukicoder.me/problems/no/3148
+template <class S> struct TreePopOrderOptimization {
+    std::vector<std::vector<int>> to;
+    std::vector<S> labels;
+    int root = -1;
+    std::vector<S> first_slope;
+    std::vector<int> par;
+
+    TreePopOrderOptimization(const std::vector<std::vector<int>> &to, const std::vector<S> &labels,
+                             int root)
+        : to(to), labels(labels), root(root), first_slope(to.size()), par(to.size(), -1) {
+
+        using Pque = std::priority_queue<S, std::vector<S>, std::greater<S>>;
+        auto rec = [&](auto &&self, int now, int prv) -> Pque {
+            std::vector<Pque> chs;
+
+            for (int nxt : to[now]) {
+                if (nxt == prv) continue;
+                assert(par[nxt] == -1);
+                par[nxt] = now;
+                chs.emplace_back(self(self, nxt, now));
+            }
+
+            const S &v = labels.at(now);
+            if (chs.empty()) {
+                first_slope[now] = v;
+                Pque pq;
+                pq.emplace(v);
+                return pq;
+            } else {
+                S first = v;
+
+                const int idx = std::max_element(chs.begin(), chs.end(),
+                                                 [](const auto &a, const auto &b) {
+                                                     return a.size() < b.size();
+                                                 }) -
+                                chs.begin();
+                std::swap(chs[idx], chs.front());
+
+                for (int i = 1; i < (int)chs.size(); ++i) {
+                    while (!chs[i].empty()) {
+                        chs.front().emplace(chs[i].top());
+                        chs[i].pop();
+                    }
+                }
+
+                while (!chs.front().empty() and chs.front().top() < first) {
+                    first += chs.front().top();
+                    chs.front().pop();
+                }
+                chs.front().emplace(first_slope[now] = first);
+                return std::move(chs.front());
+            }
+        };
+
+        rec(rec, root, -1);
+    }
+
+    std::pair<std::vector<int>, S> Solve() const { return SolveSubtree(root); }
+
+    // Generate optimal pop order of the subproblem rooted at `r`.
+    std::pair<std::vector<int>, S> SolveSubtree(int r) const {
+        using P = std::pair<S, int>;
+        std::priority_queue<P, std::vector<P>, std::greater<P>> pq;
+        pq.emplace(first_slope.at(r), r);
+
+        std::vector<int> order;
+        S ret = labels.at(r);
+        while (!pq.empty()) {
+            const int idx = pq.top().second;
+            order.emplace_back(idx);
+            pq.pop();
+            if (idx != r) ret += labels.at(idx);
+
+            for (int nxt : to.at(idx)) {
+                if (nxt == par.at(idx)) continue;
+                pq.emplace(first_slope.at(nxt), nxt);
+            }
+        }
+
+        return {order, ret};
+    }
+};
+
+template <class T> struct Vector01onTree {
+    T x, y;
+    T res;
+    Vector01onTree(T x, T y) : x(x), y(y), res(0) {}
+    Vector01onTree() : x(0), y(0), res(0) {}
+    bool operator<(const Vector01onTree &r) const {
+        if (x == 0 and y == 0) return false;
+        if (r.x == 0 and r.y == 0) return true;
+        if (x == 0 and r.x == 0) return y < r.y;
+        if (x == 0) return false;
+        if (r.x == 0) return true;
+        return y * r.x < x * r.y; // be careful of overflow
+    }
+    bool operator>(const Vector01onTree &r) const { return r < *this; }
+
+    void operator+=(const Vector01onTree &r) {
+        res += r.res + y * r.x;
+        x += r.x;
+        y += r.y;
+    }
+};
+
+template <class T>
+std::pair<std::vector<int>, T>
+Solve01OnTree(const std::vector<std::vector<int>> &to, const std::vector<T> &xs,
+              const std::vector<T> &ys, int root) {
+
+    const int n = to.size();
+    std::vector<Vector01onTree<T>> labels;
+    for (int i = 0; i < n; ++i) labels.emplace_back(xs.at(i), ys.at(i));
+
+    const TreePopOrderOptimization<Vector01onTree<T>> tpo(to, labels, root);
+    auto [order, all_prod] = tpo.Solve();
+    return {order, all_prod.res};
+}

other_algorithms/tree_pop_order_optimization.md

@@ -0,0 +1,70 @@
+---
+title: Tree pop order optimization / "01 on Tree" （木の根から 2 次元ベクトルや 01 文字列などを pop する順列に関する最小化）
+documentation_of: ./tree_pop_order_optimization.hpp
+---
+
+いわゆる "01 on Tree" を解く．
+
+## 使用方法
+
+### オーソドックスな "01 on Tree"
+
+```cpp
+int N;
+vector<vector<int>> to(N);
+
+vector<long long> x(N), y(N);
+
+const auto [order, inversions] = Solve01OnTree(to, x, y, 0);
+cout << inversions << '\n';
+for (auto e : order) cout << e << ' ';
+```
+
+### より一般的な構造
+
+以下は [28147번: Fail Fast](https://www.acmicpc.net/problem/28147) の例．まず各頂点が持つデータを表現する構造体を定義する． `operator+=` の内容に注意せよ．
+
+```cpp
+struct S {
+    double x, y;
+    S(double x, double y) : x(x), y(y) {}
+    S() : x(0), y(0) {}
+    bool operator<(const S &r) const {
+        if (x == 0 and y == 0) return false;
+        if (r.x == 0 and r.y == 0) return true;
+        if (x == 0 and r.x == 0) return y < r.y;
+        if (x == 0) return false;
+        if (r.x == 0) return true;
+        return y * r.x < x * r.y; // be careful of overflow
+    }
+    bool operator>(const S &r) const { return r < *this; }
+
+    void operator+=(const S &r) {
+        y += r.y * (1 - x);
+        x = x + (1 - x) * r.x;
+    }
+};
+```
+
+あとは，以下のように `TreePopOrderOptimization()` を呼んでやればよい．戻り値の `.first` は最適な pop 順の順列， `.second` はその順に　`S` の元の総積を（ `+=` で）とったときの値である．
+
+```cpp
+vector<vector<int>> to(N);
+vector<S> vals(N);
+
+auto [seq, all_prod] = TreePopOrderOptimization(to, vals, 0).Solve();
+```
+
+## 問題例
+
+- [Library Checker: Rooted Tree Topological Order with Minimum Inversions](https://judge.yosupo.jp/problem/rooted_tree_topological_order_with_minimum_inversions)
+- [AtCoder Grand Contest 023 F - 01 on Tree](https://atcoder.jp/contests/agc023/tasks/agc023_f)
+- [AtCoder Beginner Contest 376 G - Treasure Hunting](https://atcoder.jp/contests/abc376/tasks/abc376_g)
+- [No.3148 Min-Cost Destruction of Parentheses - yukicoder](https://yukicoder.me/problems/no/3148)
+- [28147번: Fail Fast](https://www.acmicpc.net/problem/28147)
+  - 通常の 01 on Tree とは異なる演算の入ったデータ構造を載せるタイプの問題．
+
+## リンク
+
+- [解説 - AtCoder Beginner Contest 376（Promotion of AtCoder Career Design DAY）](https://atcoder.jp/contests/abc376/editorial/11196)
+- [241203_01 on Tree](https://acompany-ac.notion.site/241203_01-on-Tree-151269d8558680b2b639d7bfcbff2b20)