verify/standalone-01-on-tree.test.cpp

• View this file on GitHub
• Last update: 2026-07-11 16:54:42+09:00

Depends on

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

Code

// competitive-verifier: STANDALONE

#include <algorithm>
#include <cassert>
#include <limits>
#include <utility>
#include <vector>

#include "../graph/tree/01-on-tree.hpp"

long long brute_01_on_tree(int n, const std::vector<std::pair<int, int>> &edges,
                           const std::vector<long long> &c0,
                           const std::vector<long long> &c1, int root) {
    std::vector<std::vector<int>> graph(n);
    for (const auto &[u, v] : edges) {
        graph[u].push_back(v);
        graph[v].push_back(u);
    }

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

    std::vector<int> order(n), position(n);
    for (int v = 0; v < n; ++v) {
        order[v] = v;
    }

    long long answer = std::numeric_limits<long long>::max();
    do {
        for (int i = 0; i < n; ++i) {
            position[order[i]] = i;
        }
        bool valid = true;
        for (int v = 0; v < n; ++v) {
            if (parent[v] != -1 && position[parent[v]] > position[v]) {
                valid = false;
            }
        }
        if (!valid) {
            continue;
        }

        long long current = 0;
        for (int i = 0; i < n; ++i) {
            for (int j = i + 1; j < n; ++j) {
                current += c1[order[i]] * c0[order[j]];
            }
        }
        answer = std::min(answer, current);
    } while (std::next_permutation(order.begin(), order.end()));
    return answer;
}

void test_01_on_tree(int n, const std::vector<std::pair<int, int>> &edges,
                     const std::vector<long long> &c0,
                     const std::vector<long long> &c1, int root) {
    const long long expected = brute_01_on_tree(n, edges, c0, c1, root);
    const long long actual = solve_01_on_tree(n, edges, c0, c1, root);
    assert(actual == expected);
}

void test_all_small_counts(int n,
                           const std::vector<std::pair<int, int>> &edges) {
    int pattern_count = 1;
    for (int i = 0; i < 2 * n; ++i) {
        pattern_count *= 3;
    }
    for (int code = 0; code < pattern_count; ++code) {
        int current_code = code;
        std::vector<long long> c0(n), c1(n);
        for (int v = 0; v < n; ++v) {
            c0[v] = current_code % 3;
            current_code /= 3;
            c1[v] = current_code % 3;
            current_code /= 3;
        }
        test_01_on_tree(n, edges, c0, c1, 0);
    }
}

void enumerate_parent_trees(int n, int v,
                            std::vector<std::pair<int, int>> &edges) {
    if (v == n) {
        test_all_small_counts(n, edges);
        return;
    }
    for (int parent = 0; parent < v; ++parent) {
        edges.push_back({parent, v});
        enumerate_parent_trees(n, v + 1, edges);
        edges.pop_back();
    }
}

int main() {
    for (int n = 1; n <= 4; ++n) {
        std::vector<std::pair<int, int>> edges;
        enumerate_parent_trees(n, 1, edges);
    }

    const std::vector<std::pair<int, int>> edges{
        {0, 1}, {0, 2}, {2, 3}, {2, 4}};
    const std::vector<long long> c0{0, 3, 0, 2, 1};
    const std::vector<long long> c1{0, 0, 4, 1, 0};
    for (int root = 0; root < 5; ++root) {
        test_01_on_tree(5, edges, c0, c1, root);
    }

    const std::vector<unsigned long long> large_c0{0, 4'000'000'000ULL};
    const std::vector<unsigned long long> large_c1{4'000'000'000ULL, 0};
    const auto large_answer = solve_01_on_tree(2, {{0, 1}}, large_c0, large_c1);
    assert(large_answer == 16'000'000'000'000'000'000ULL);

    return 0;
}

#line 1 "verify/standalone-01-on-tree.test.cpp"
// competitive-verifier: STANDALONE

#include <algorithm>
#include <cassert>
#include <limits>
#include <utility>
#include <vector>

#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) である。

#line 10 "graph/tree/01-on-tree.hpp"
#include <queue>
#include <type_traits>
#line 14 "graph/tree/01-on-tree.hpp"

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;
    [[maybe_unused]] int visited = 0;
    std::vector<int> stack;
    stack.reserve(n);
    stack.push_back(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.zero == b.zero && a.one == b.one) {
                return a.vertex < b.vertex;
            }
            if (a.one == Count{} && b.one == Count{}) {
                return a.vertex < b.vertex;
            }
            if (a.one == Count{}) {
                return false;
            }
            if (b.one == Count{}) {
                return true;
            }
            if (a.zero == Count{} && b.zero == Count{}) {
                return a.vertex < b.vertex;
            }
            if (a.zero == Count{}) {
                return true;
            }
            if (b.zero == Count{}) {
                return false;
            }
            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 = std::move(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::vector<QueueNode> queue_storage;
    queue_storage.reserve(n);
    std::priority_queue<QueueNode, std::vector<QueueNode>, QueueCompare> que(
        QueueCompare{}, std::move(queue_storage));
    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) {
        int v = -1;
        while (true) {
            const QueueNode &current = que.top();
            const int current_vertex = current.vertex;
            const int current_version = current.version;
            que.pop();
            v = find(current_vertex);
            if (v == current_vertex && current_version == version[v]) {
                break;
            }
        }

        const int p = find(up[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;
}


#line 10 "verify/standalone-01-on-tree.test.cpp"

long long brute_01_on_tree(int n, const std::vector<std::pair<int, int>> &edges,
                           const std::vector<long long> &c0,
                           const std::vector<long long> &c1, int root) {
    std::vector<std::vector<int>> graph(n);
    for (const auto &[u, v] : edges) {
        graph[u].push_back(v);
        graph[v].push_back(u);
    }

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

    std::vector<int> order(n), position(n);
    for (int v = 0; v < n; ++v) {
        order[v] = v;
    }

    long long answer = std::numeric_limits<long long>::max();
    do {
        for (int i = 0; i < n; ++i) {
            position[order[i]] = i;
        }
        bool valid = true;
        for (int v = 0; v < n; ++v) {
            if (parent[v] != -1 && position[parent[v]] > position[v]) {
                valid = false;
            }
        }
        if (!valid) {
            continue;
        }

        long long current = 0;
        for (int i = 0; i < n; ++i) {
            for (int j = i + 1; j < n; ++j) {
                current += c1[order[i]] * c0[order[j]];
            }
        }
        answer = std::min(answer, current);
    } while (std::next_permutation(order.begin(), order.end()));
    return answer;
}

void test_01_on_tree(int n, const std::vector<std::pair<int, int>> &edges,
                     const std::vector<long long> &c0,
                     const std::vector<long long> &c1, int root) {
    const long long expected = brute_01_on_tree(n, edges, c0, c1, root);
    const long long actual = solve_01_on_tree(n, edges, c0, c1, root);
    assert(actual == expected);
}

void test_all_small_counts(int n,
                           const std::vector<std::pair<int, int>> &edges) {
    int pattern_count = 1;
    for (int i = 0; i < 2 * n; ++i) {
        pattern_count *= 3;
    }
    for (int code = 0; code < pattern_count; ++code) {
        int current_code = code;
        std::vector<long long> c0(n), c1(n);
        for (int v = 0; v < n; ++v) {
            c0[v] = current_code % 3;
            current_code /= 3;
            c1[v] = current_code % 3;
            current_code /= 3;
        }
        test_01_on_tree(n, edges, c0, c1, 0);
    }
}

void enumerate_parent_trees(int n, int v,
                            std::vector<std::pair<int, int>> &edges) {
    if (v == n) {
        test_all_small_counts(n, edges);
        return;
    }
    for (int parent = 0; parent < v; ++parent) {
        edges.push_back({parent, v});
        enumerate_parent_trees(n, v + 1, edges);
        edges.pop_back();
    }
}

int main() {
    for (int n = 1; n <= 4; ++n) {
        std::vector<std::pair<int, int>> edges;
        enumerate_parent_trees(n, 1, edges);
    }

    const std::vector<std::pair<int, int>> edges{
        {0, 1}, {0, 2}, {2, 3}, {2, 4}};
    const std::vector<long long> c0{0, 3, 0, 2, 1};
    const std::vector<long long> c1{0, 0, 4, 1, 0};
    for (int root = 0; root < 5; ++root) {
        test_01_on_tree(5, edges, c0, c1, root);
    }

    const std::vector<unsigned long long> large_c0{0, 4'000'000'000ULL};
    const std::vector<unsigned long long> large_c1{4'000'000'000ULL, 0};
    const auto large_answer = solve_01_on_tree(2, {{0, 1}}, large_c0, large_c1);
    assert(large_answer == 16'000'000'000'000'000'000ULL);

    return 0;
}

Back to top page