Tree pop order optimization / "01 on Tree" （木の根から 2 次元ベクトルや 01 文字列などを pop する順列に関する最小化）
(other_algorithms/tree_pop_order_optimization.hpp)

View this file on GitHub
Last update: 2025-08-05 23:34:24+09:00
Include: #include "other_algorithms/tree_pop_order_optimization.hpp"
Link: https://judge.yosupo.jp/problem/rooted_tree_topological_order_with_minimum_inversions
Link: https://yukicoder.me/problems/no/3148

いわゆる “01 on Tree” を解く．

使用方法

オーソドックスな “01 on Tree”

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 の例．まず各頂点が持つデータを表現する構造体を定義する． operator+= の内容に注意せよ．

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 の元の総積を（ += で）とったときの値である．

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
AtCoder Grand Contest 023 F - 01 on Tree
AtCoder Beginner Contest 376 G - Treasure Hunting
No.3148 Min-Cost Destruction of Parentheses - yukicoder
28147번: Fail Fast
- 通常の 01 on Tree とは異なる演算の入ったデータ構造を載せるタイプの問題．

リンク

Verified with

Code

#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};
}

#line 2 "other_algorithms/tree_pop_order_optimization.hpp"
#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};
}

Tree pop order optimization / "01 on Tree" （木の根から 2 次元ベクトルや 01 文字列などを pop する順列に関する最小化） (other_algorithms/tree_pop_order_optimization.hpp)