01 on Tree (graph/tree/01-on-tree.hpp)

• View this file on GitHub
• View document part on GitHub
• Last update: 2026-05-31 17:59:37+09:00
• Include: #include "graph/tree/01-on-tree.hpp"

概要

• 根付き木の各頂点 $v$ に $V_v=(0,\ldots,0,1,\ldots,1)$ を置く。
• $0$ の個数は $c^{(0)}_v$ 個、$1$ の個数は $c^{(1)}_v$ 個である。
• 親が子より左に出るように頂点を並べるとき、$V_v$ をその順に連結してできる列の転倒数の最小値を求める。
• 辺は無向辺の両端で与え、根を指定して親子関係を定める。
• 根でない成分のうち $c^{(0)}/c^{(1)}$ が最大のものを親へ縮約する。$c^{(1)}=0$ の成分は最大として扱う。
• 比の比較は積を作らずに行う。

使い方

• solve_01_on_tree<T>(n, edges, c0, c1, root = 0)
  • 頂点数 $n$ の根付き木について、転倒数の最小値を T 型で返す。
  • T を省略した場合は、個数型を基にした型で返す。個数型が long long より狭い場合は long long または unsigned long long に上げる。
  • T は返り値の型であり、順序比較できない型でもよい。
  • T が個数型より広くプリミティブで、かつ浮動小数点でない型である場合、内部の累積個数型にも T を使う。
  • c0, c1 の要素型を入力の個数型とする。個数型は非負整数型で、順序比較、等値比較、加算、除算、剰余、返り値型への変換ができる必要がある。
  • T を modint にする場合も、c0, c1 は整数型の配列で渡す。
  • 頂点番号は 0-based indexing とする。
  • edges は木の無向辺の両端の列である。辺の向きは参照しない。
  • c0[v] は頂点 $v$ の列に含まれる $0$ の個数である。
  • c1[v] は頂点 $v$ の列に含まれる $1$ の個数である。
  • 頂点 $v$ の列は $c^{(0)}_v$ 個の $0$ の後に $c^{(1)}_v$ 個の $1$ を置いたものとする。
  • root は根の頂点番号である。$0\le \mathrm{root}<n$ である必要がある。
  • 前提: $n\ge 1$、辺数は $n-1$、c0 と c1 の長さは $n$ である必要がある。
  • 前提: edges は $1$ つの木を表す必要がある。頂点番号が範囲外、自己ループ、閉路、到達不能な頂点がある場合は assert に失敗する。
  • 前提: 各 $c^{(0)}_v,c^{(1)}_v$ は非負の個数である必要がある。
  • 前提: 内部の累積個数と返り値型 T 上の答えが用途に対して十分な範囲を持つ必要がある。
  • 比の比較で積は作らないため、比較用の積が個数型に収まる必要はない。

計算量

• 頂点数を $n$ とする。
• 固定幅整数の算術演算を $O(1)$ として、$O(n\log n)$ 時間、$O(n)$ 空間である。

Verified with

• verify/yosupo-rooted-tree-topological-order-with-minimum-inversions.test.cpp
• verify/yukicoder-3148.test.cpp

Code

#ifndef GRAPH_TREE_01_ON_TREE_HPP
#define GRAPH_TREE_01_ON_TREE_HPP

// 根付き木上で、親が子より左に出る順序の 01 列の転倒数最小値を求める。
// 頂点 v には c0[v] 個の 0 の後に c1[v] 個の 1 を置いた列が書かれている。
// 辺は無向辺の両端で与え、root を根として親子関係を定める。
// c0[v], c1[v] は非負の個数である。計算量は O(n log n) である。

#include <cassert>
#include <queue>
#include <type_traits>
#include <utility>
#include <vector>

namespace zero_one_on_tree_impl {
template <class InputCount, class Preferred>
using count_type_t = std::conditional_t<
    !std::is_class_v<Preferred> && !std::is_floating_point_v<Preferred> &&
        (sizeof(InputCount) < sizeof(Preferred)),
    Preferred,
    std::conditional_t<(sizeof(InputCount) < sizeof(long long)),
                       std::conditional_t<std::is_signed_v<InputCount>,
                                          long long, unsigned long long>,
                       InputCount>>;

template <class Count>
int compare_fraction(Count a_num, Count a_den, Count b_num, Count b_den) {
    bool reversed = false;
    while (true) {
        const Count a_quot = a_num / a_den;
        const Count b_quot = b_num / b_den;
        if (a_quot != b_quot) {
            const int result = a_quot < b_quot ? -1 : 1;
            return reversed ? -result : result;
        }

        a_num %= a_den;
        b_num %= b_den;
        if (a_num == Count{} || b_num == Count{}) {
            int result = 0;
            if (a_num != b_num) {
                result = a_num == Count{} ? -1 : 1;
            }
            return reversed ? -result : result;
        }

        std::swap(a_num, a_den);
        std::swap(b_num, b_den);
        reversed = !reversed;
    }
}
} // namespace zero_one_on_tree_impl

template <class T = void, class InputCount>
std::conditional_t<std::is_void_v<T>,
                   zero_one_on_tree_impl::count_type_t<InputCount, InputCount>,
                   T>
solve_01_on_tree(int n, const std::vector<std::pair<int, int>> &edges,
                 const std::vector<InputCount> &c0,
                 const std::vector<InputCount> &c1, int root = 0) {
    using Preferred = std::conditional_t<std::is_void_v<T>, InputCount, T>;
    using Count = zero_one_on_tree_impl::count_type_t<InputCount, Preferred>;
    using Answer = std::conditional_t<std::is_void_v<T>, Count, T>;

    assert(n >= 1);
    assert(static_cast<int>(edges.size()) == n - 1);
    assert(static_cast<int>(c0.size()) == n);
    assert(static_cast<int>(c1.size()) == n);
    assert(0 <= root && root < n);

    std::vector<std::vector<int>> graph(n);
    std::vector<Count> zero(n), one(n);
    for (int v = 0; v < n; ++v) {
        assert(!(c0[v] < InputCount{}));
        assert(!(c1[v] < InputCount{}));
        zero[v] = static_cast<Count>(c0[v]);
        one[v] = static_cast<Count>(c1[v]);
    }
    for (const auto &[from, to] : edges) {
        assert(0 <= from && from < n);
        assert(0 <= to && to < n);
        assert(from != to);
        graph[from].push_back(to);
        graph[to].push_back(from);
    }

    std::vector<int> parent(n, -2);
    parent[root] = -1;
    int visited = 0;
    std::vector<int> stack{root};
    while (!stack.empty()) {
        const int v = stack.back();
        stack.pop_back();
        ++visited;
        for (int to : graph[v]) {
            if (to == parent[v]) {
                continue;
            }
            assert(parent[to] == -2);
            parent[to] = v;
            stack.push_back(to);
        }
    }
    assert(visited == n);

    struct QueueNode {
        int vertex;
        int version;
        Count zero;
        Count one;
    };

    struct QueueCompare {
        bool operator()(const QueueNode &a, const QueueNode &b) const {
            if (a.one == Count{} && b.one == Count{}) {
                return a.vertex < b.vertex;
            }
            if (a.one == Count{}) {
                return false;
            }
            if (b.one == Count{}) {
                return true;
            }
            const int result = zero_one_on_tree_impl::compare_fraction(
                a.zero, a.one, b.zero, b.one);
            if (result < 0) {
                return true;
            }
            if (result > 0) {
                return false;
            }
            return a.vertex < b.vertex;
        }
    };

    std::vector<int> leader(n), up = parent, version(n, 0);
    for (int v = 0; v < n; ++v) {
        leader[v] = v;
    }

    auto find = [&](int v) {
        while (leader[v] != v) {
            leader[v] = leader[leader[v]];
            v = leader[v];
        }
        return v;
    };

    std::priority_queue<QueueNode, std::vector<QueueNode>, QueueCompare> que;
    auto push = [&](int v) {
        if (up[v] != -1) {
            que.push(QueueNode{v, version[v], zero[v], one[v]});
        }
    };

    for (int v = 0; v < n; ++v) {
        push(v);
    }

    Answer answer{};
    for (int merge_count = 0; merge_count < n - 1; ++merge_count) {
        QueueNode current{};
        while (true) {
            assert(!que.empty());
            current = que.top();
            que.pop();
            const int v = find(current.vertex);
            if (v == current.vertex && current.version == version[v]) {
                break;
            }
        }

        const int v = current.vertex;
        const int p = find(up[v]);
        assert(p != v);

        answer += static_cast<Answer>(one[p]) * static_cast<Answer>(zero[v]);
        zero[p] += zero[v];
        one[p] += one[v];
        leader[v] = p;
        ++version[p];
        push(p);
    }

    return answer;
}

#endif

#line 1 "graph/tree/01-on-tree.hpp"



// 根付き木上で、親が子より左に出る順序の 01 列の転倒数最小値を求める。
// 頂点 v には c0[v] 個の 0 の後に c1[v] 個の 1 を置いた列が書かれている。
// 辺は無向辺の両端で与え、root を根として親子関係を定める。
// c0[v], c1[v] は非負の個数である。計算量は O(n log n) である。

#include <cassert>
#include <queue>
#include <type_traits>
#include <utility>
#include <vector>

namespace zero_one_on_tree_impl {
template <class InputCount, class Preferred>
using count_type_t = std::conditional_t<
    !std::is_class_v<Preferred> && !std::is_floating_point_v<Preferred> &&
        (sizeof(InputCount) < sizeof(Preferred)),
    Preferred,
    std::conditional_t<(sizeof(InputCount) < sizeof(long long)),
                       std::conditional_t<std::is_signed_v<InputCount>,
                                          long long, unsigned long long>,
                       InputCount>>;

template <class Count>
int compare_fraction(Count a_num, Count a_den, Count b_num, Count b_den) {
    bool reversed = false;
    while (true) {
        const Count a_quot = a_num / a_den;
        const Count b_quot = b_num / b_den;
        if (a_quot != b_quot) {
            const int result = a_quot < b_quot ? -1 : 1;
            return reversed ? -result : result;
        }

        a_num %= a_den;
        b_num %= b_den;
        if (a_num == Count{} || b_num == Count{}) {
            int result = 0;
            if (a_num != b_num) {
                result = a_num == Count{} ? -1 : 1;
            }
            return reversed ? -result : result;
        }

        std::swap(a_num, a_den);
        std::swap(b_num, b_den);
        reversed = !reversed;
    }
}
} // namespace zero_one_on_tree_impl

template <class T = void, class InputCount>
std::conditional_t<std::is_void_v<T>,
                   zero_one_on_tree_impl::count_type_t<InputCount, InputCount>,
                   T>
solve_01_on_tree(int n, const std::vector<std::pair<int, int>> &edges,
                 const std::vector<InputCount> &c0,
                 const std::vector<InputCount> &c1, int root = 0) {
    using Preferred = std::conditional_t<std::is_void_v<T>, InputCount, T>;
    using Count = zero_one_on_tree_impl::count_type_t<InputCount, Preferred>;
    using Answer = std::conditional_t<std::is_void_v<T>, Count, T>;

    assert(n >= 1);
    assert(static_cast<int>(edges.size()) == n - 1);
    assert(static_cast<int>(c0.size()) == n);
    assert(static_cast<int>(c1.size()) == n);
    assert(0 <= root && root < n);

    std::vector<std::vector<int>> graph(n);
    std::vector<Count> zero(n), one(n);
    for (int v = 0; v < n; ++v) {
        assert(!(c0[v] < InputCount{}));
        assert(!(c1[v] < InputCount{}));
        zero[v] = static_cast<Count>(c0[v]);
        one[v] = static_cast<Count>(c1[v]);
    }
    for (const auto &[from, to] : edges) {
        assert(0 <= from && from < n);
        assert(0 <= to && to < n);
        assert(from != to);
        graph[from].push_back(to);
        graph[to].push_back(from);
    }

    std::vector<int> parent(n, -2);
    parent[root] = -1;
    int visited = 0;
    std::vector<int> stack{root};
    while (!stack.empty()) {
        const int v = stack.back();
        stack.pop_back();
        ++visited;
        for (int to : graph[v]) {
            if (to == parent[v]) {
                continue;
            }
            assert(parent[to] == -2);
            parent[to] = v;
            stack.push_back(to);
        }
    }
    assert(visited == n);

    struct QueueNode {
        int vertex;
        int version;
        Count zero;
        Count one;
    };

    struct QueueCompare {
        bool operator()(const QueueNode &a, const QueueNode &b) const {
            if (a.one == Count{} && b.one == Count{}) {
                return a.vertex < b.vertex;
            }
            if (a.one == Count{}) {
                return false;
            }
            if (b.one == Count{}) {
                return true;
            }
            const int result = zero_one_on_tree_impl::compare_fraction(
                a.zero, a.one, b.zero, b.one);
            if (result < 0) {
                return true;
            }
            if (result > 0) {
                return false;
            }
            return a.vertex < b.vertex;
        }
    };

    std::vector<int> leader(n), up = parent, version(n, 0);
    for (int v = 0; v < n; ++v) {
        leader[v] = v;
    }

    auto find = [&](int v) {
        while (leader[v] != v) {
            leader[v] = leader[leader[v]];
            v = leader[v];
        }
        return v;
    };

    std::priority_queue<QueueNode, std::vector<QueueNode>, QueueCompare> que;
    auto push = [&](int v) {
        if (up[v] != -1) {
            que.push(QueueNode{v, version[v], zero[v], one[v]});
        }
    };

    for (int v = 0; v < n; ++v) {
        push(v);
    }

    Answer answer{};
    for (int merge_count = 0; merge_count < n - 1; ++merge_count) {
        QueueNode current{};
        while (true) {
            assert(!que.empty());
            current = que.top();
            que.pop();
            const int v = find(current.vertex);
            if (v == current.vertex && current.version == version[v]) {
                break;
            }
        }

        const int v = current.vertex;
        const int p = find(up[v]);
        assert(p != v);

        answer += static_cast<Answer>(one[p]) * static_cast<Answer>(zero[v]);
        zero[p] += zero[v];
        one[p] += one[v];
        leader[v] = p;
        ++version[p];
        push(p);
    }

    return answer;
}

Back to top page