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

View this file on GitHub
Last update: 2026-08-13 01:41:40+09:00
Include: #include "graph/tree/zero_one_on_tree.hpp"

Overview

Each vertex of a rooted tree is labeled 0 or 1. Consider every ordering of the vertices in which each parent appears before its children. This header computes the minimum possible number of inversions in the resulting binary sequence.

The algorithm maintains contracted vertex groups. A non-root group with the largest ratio

[ \frac{\text{number of zeros}}{\text{number of ones}} ]

is appended to its current parent group. A group containing no ones has infinite ratio. A set maintains the ratios, and a disjoint-set structure finds the current parent group after contractions.

Public Interface

long long zero_one_on_tree(
    const std::vector<int>& parent,
    const std::vector<int>& value);

template <class T>
long long zero_one_on_tree(
    const m1une::graph::Graph<T>& graph,
    const std::vector<int>& value,
    int root = 0);

The parent-array overload uses zero-based vertices. Exactly one vertex must have parent -1; every other entry must be a valid parent and the resulting directed edges must form a tree. Each entry of value must be 0 or 1.

The graph overload expects a connected undirected tree built with add_edge. Inactive edges are ignored. For an empty graph it returns 0 without using root.

Complexity

Function	Time	Extra memory
`zero_one_on_tree(parent, value)`	$O(N \log N)$	$O(N)$
`zero_one_on_tree(graph, value, root)`	$O(N \log N)$	$O(N)$

The answer can be quadratic in N, so the return type is long long. Neither overload mutates its arguments.

Example

#include "graph/tree/zero_one_on_tree.hpp"
#include <iostream>
#include <vector>

int main() {
    std::vector<int> parent = {-1, 0, 0, 1, 2, 2};
    std::vector<int> value = {0, 1, 1, 0, 0, 0};

    std::cout << m1une::tree::zero_one_on_tree(parent, value) << '\n';
}

Depends on

Required by

Verified with

Code

#ifndef M1UNE_TREE_ZERO_ONE_ON_TREE_HPP
#define M1UNE_TREE_ZERO_ONE_ON_TREE_HPP 1

#include <cassert>
#include <set>
#include <vector>

#include "rooted_tree.hpp"

namespace m1une {
namespace tree {

inline long long zero_one_on_tree(const std::vector<int>& parent,
                                  const std::vector<int>& value) {
    const int n = int(parent.size());
    assert(int(value.size()) == n);
    if (n == 0) return 0;

    int root = -1;
    std::vector<std::vector<int>> children(n);
    for (int v = 0; v < n; v++) {
        assert(value[v] == 0 || value[v] == 1);
        if (parent[v] == -1) {
            assert(root == -1);
            root = v;
        } else {
            assert(0 <= parent[v] && parent[v] < n && parent[v] != v);
            children[parent[v]].push_back(v);
        }
    }
    assert(root != -1);

    std::vector<int> stack(1, root);
    std::vector<char> visited(n, false);
    visited[root] = true;
    int visited_count = 0;
    while (!stack.empty()) {
        const int v = stack.back();
        stack.pop_back();
        visited_count++;
        for (int child : children[v]) {
            assert(!visited[child]);
            visited[child] = true;
            stack.push_back(child);
        }
    }
    assert(visited_count == n);

    struct Component {
        long long zeros;
        long long ones;
        int vertex;
    };
    struct Compare {
        bool operator()(const Component& lhs, const Component& rhs) const {
            const long long lhs_product = lhs.zeros * rhs.ones;
            const long long rhs_product = rhs.zeros * lhs.ones;
            if (lhs_product != rhs_product) return lhs_product < rhs_product;
            return lhs.vertex < rhs.vertex;
        }
    };

    std::vector<long long> zeros(n), ones(n);
    std::vector<int> dsu(n);
    std::set<Component, Compare> components;
    for (int v = 0; v < n; v++) {
        zeros[v] = value[v] == 0;
        ones[v] = value[v] == 1;
        dsu[v] = v;
        if (v != root) components.insert(Component{zeros[v], ones[v], v});
    }

    auto leader = [&](int v) {
        int result = v;
        while (dsu[result] != result) result = dsu[result];
        while (dsu[v] != v) {
            const int next = dsu[v];
            dsu[v] = result;
            v = next;
        }
        return result;
    };

    long long answer = 0;
    while (!components.empty()) {
        auto it = components.end();
        --it;
        const Component child = *it;
        components.erase(it);

        const int p = leader(parent[child.vertex]);
        if (p != root) {
            const int erased = int(components.erase(Component{zeros[p], ones[p], p}));
            assert(erased == 1);
        }

        answer += ones[p] * zeros[child.vertex];
        zeros[p] += zeros[child.vertex];
        ones[p] += ones[child.vertex];
        dsu[child.vertex] = p;

        if (p != root) components.insert(Component{zeros[p], ones[p], p});
    }
    return answer;
}

template <class T>
long long zero_one_on_tree(const m1une::graph::Graph<T>& graph,
                           const std::vector<int>& value, int root = 0) {
    const int n = graph.size();
    assert(int(value.size()) == n);
    if (n == 0) return 0;
    assert(0 <= root && root < n);
    assert(int(graph.edges().size()) == n - 1);

    RootedTree<T> rooted_tree(graph, root);
    assert(int(rooted_tree.order.size()) == n);
    return zero_one_on_tree(rooted_tree.parent, value);
}

}  // namespace tree
}  // namespace m1une

#endif  // M1UNE_TREE_ZERO_ONE_ON_TREE_HPP

#line 1 "graph/tree/zero_one_on_tree.hpp"



#include <cassert>
#include <set>
#include <vector>

#line 1 "graph/tree/rooted_tree.hpp"



#include <algorithm>
#line 7 "graph/tree/rooted_tree.hpp"

#line 1 "graph/graph.hpp"



#include <array>
#line 6 "graph/graph.hpp"
#include <utility>
#line 8 "graph/graph.hpp"

namespace m1une {
namespace graph {

template <class T = int>
struct Edge {
    using cost_type = T;

    int from;
    int to;
    T cost;
    int id;
    bool alive;

    Edge() : from(-1), to(-1), cost(T()), id(-1), alive(true) {}
    Edge(int from_, int to_, T cost_ = T(1), int id_ = -1, bool alive_ = true)
        : from(from_), to(to_), cost(cost_), id(id_), alive(alive_) {}

    int other(int v) const {
        assert(v == from || v == to);
        return from ^ to ^ v;
    }
};

template <class T = int>
struct Graph {
    using edge_type = Edge<T>;
    using cost_type = T;

   private:
    struct EdgePositions {
        std::array<std::pair<int, int>, 2> value{};
        int size = 0;

        void push_back(std::pair<int, int> position) {
            assert(size < 2);
            value[size++] = position;
        }
    };

    int _n;
    int _edge_count;
    std::vector<std::vector<edge_type>> _g;
    std::vector<EdgePositions> _edge_positions;

   public:
    Graph() : _n(0), _edge_count(0) {}
    explicit Graph(int n) : _n(n), _edge_count(0), _g(n) {
        assert(0 <= n);
    }

    int size() const {
        return _n;
    }

    bool empty() const {
        return _n == 0;
    }

    int edge_count() const {
        return _edge_count;
    }

    int add_vertex() {
        _g.emplace_back();
        return _n++;
    }

    int add_directed_edge(int from, int to, T cost = T(1)) {
        assert(0 <= from && from < _n);
        assert(0 <= to && to < _n);
        int id = _edge_count++;
        int idx = int(_g[from].size());
        _g[from].push_back(edge_type(from, to, cost, id));
        _edge_positions.emplace_back();
        _edge_positions.back().push_back({from, idx});
        return id;
    }

    int add_edge(int u, int v, T cost = T(1)) {
        assert(0 <= u && u < _n);
        assert(0 <= v && v < _n);
        int id = _edge_count++;
        int u_idx = int(_g[u].size());
        _g[u].push_back(edge_type(u, v, cost, id));
        int v_idx = int(_g[v].size());
        _g[v].push_back(edge_type(v, u, cost, id));
        _edge_positions.emplace_back();
        _edge_positions.back().push_back({u, u_idx});
        _edge_positions.back().push_back({v, v_idx});
        return id;
    }

    void set_edge_alive(int id, bool alive) {
        assert(0 <= id && id < _edge_count);
        for (int i = 0; i < _edge_positions[id].size; ++i) {
            auto [v, idx] = _edge_positions[id].value[i];
            _g[v][idx].alive = alive;
        }
    }

    void erase_edge(int id) {
        set_edge_alive(id, false);
    }

    void revive_edge(int id) {
        set_edge_alive(id, true);
    }

    bool is_edge_alive(int id) const {
        assert(0 <= id && id < _edge_count);
        assert(_edge_positions[id].size != 0);
        auto [v, idx] = _edge_positions[id].value[0];
        return _g[v][idx].alive;
    }

    const std::vector<edge_type>& operator[](int v) const {
        assert(0 <= v && v < _n);
        return _g[v];
    }

    std::vector<edge_type>& operator[](int v) {
        assert(0 <= v && v < _n);
        return _g[v];
    }

    const std::vector<std::vector<edge_type>>& adjacency() const {
        return _g;
    }

    std::vector<std::vector<edge_type>>& adjacency() {
        return _g;
    }

    std::vector<edge_type> edges(bool include_inactive = false) const {
        std::vector<edge_type> result;
        result.reserve(_edge_count);
        std::vector<char> used(_edge_count, false);
        for (int v = 0; v < _n; v++) {
            for (const auto& e : _g[v]) {
                if (!include_inactive && !e.alive) continue;
                if (0 <= e.id && e.id < _edge_count) {
                    if (used[e.id]) continue;
                    used[e.id] = true;
                }
                result.push_back(e);
            }
        }
        return result;
    }

    Graph reversed() const {
        Graph result(_n);
        result._edge_count = _edge_count;
        result._edge_positions.assign(_edge_count, {});
        for (int v = 0; v < _n; v++) {
            for (const auto& e : _g[v]) {
                int idx = int(result._g[e.to].size());
                result._g[e.to].push_back(edge_type(e.to, e.from, e.cost, e.id, e.alive));
                if (0 <= e.id && e.id < _edge_count) result._edge_positions[e.id].push_back({e.to, idx});
            }
        }
        return result;
    }
};

}  // namespace graph
}  // namespace m1une


#line 9 "graph/tree/rooted_tree.hpp"

namespace m1une {
namespace tree {

template <class T = int>
struct RootedTree {
    using cost_type = T;
    using edge_type = m1une::graph::Edge<T>;

    int root;
    std::vector<int> parent;
    std::vector<int> parent_edge;
    std::vector<int> depth;
    std::vector<T> dist;
    std::vector<int> subtree_size;
    std::vector<int> tin;
    std::vector<int> tout;
    std::vector<int> order;
    std::vector<std::vector<int>> up;

   private:
    int _n;
    int _log;

    void check_vertex(int v) const {
        assert(0 <= v && v < _n);
        assert(tin[v] != -1);
    }

   public:
    RootedTree() : root(-1), _n(0), _log(0) {}
    explicit RootedTree(const m1une::graph::Graph<T>& g, int root_ = 0) {
        build(g, root_);
    }

    void build(const m1une::graph::Graph<T>& g, int root_ = 0) {
        _n = g.size();
        root = _n == 0 ? -1 : root_;
        _log = 1;
        while ((1U << _log) <= (unsigned int)(std::max(1, _n))) _log++;

        parent.assign(_n, -1);
        parent_edge.assign(_n, -1);
        depth.assign(_n, 0);
        dist.assign(_n, T(0));
        subtree_size.assign(_n, 0);
        tin.assign(_n, -1);
        tout.assign(_n, -1);
        order.clear();
        order.reserve(_n);
        up.assign(_log, std::vector<int>(_n, -1));

        if (_n == 0) return;
        assert(0 <= root && root < _n);

        struct Frame {
            int v;
            int state;
        };

        std::vector<char> visited(_n, false);
        std::vector<Frame> stack;
        stack.push_back({root, 0});
        visited[root] = true;
        int timer = 0;

        while (!stack.empty()) {
            Frame frame = stack.back();
            stack.pop_back();
            int v = frame.v;
            if (frame.state == 0) {
                tin[v] = timer++;
                order.push_back(v);
                up[0][v] = parent[v];
                for (int k = 1; k < _log; k++) {
                    int p = up[k - 1][v];
                    up[k][v] = p == -1 ? -1 : up[k - 1][p];
                }

                stack.push_back({v, 1});
                const auto& adj = g[v];
                for (int i = int(adj.size()) - 1; i >= 0; i--) {
                    const auto& e = adj[i];
                    if (!e.alive) continue;
                    if (visited[e.to]) continue;
                    visited[e.to] = true;
                    parent[e.to] = v;
                    parent_edge[e.to] = e.id;
                    depth[e.to] = depth[v] + 1;
                    dist[e.to] = dist[v] + e.cost;
                    stack.push_back({e.to, 0});
                }
            } else {
                subtree_size[v]++;
                if (parent[v] != -1) subtree_size[parent[v]] += subtree_size[v];
                tout[v] = timer;
            }
        }
    }

    int size() const {
        return _n;
    }

    bool empty() const {
        return _n == 0;
    }

    int log() const {
        return _log;
    }

    bool is_ancestor(int u, int v) const {
        check_vertex(u);
        check_vertex(v);
        return tin[u] <= tin[v] && tout[v] <= tout[u];
    }

    bool in_subtree(int v, int u) const {
        return is_ancestor(u, v);
    }

    int kth_ancestor(int v, int k) const {
        check_vertex(v);
        assert(0 <= k);
        int bit = 0;
        while (k > 0 && v != -1) {
            if (k & 1) {
                if (_log <= bit) return -1;
                v = up[bit][v];
            }
            k >>= 1;
            bit++;
        }
        return v;
    }

    int lca(int u, int v) const {
        check_vertex(u);
        check_vertex(v);
        if (depth[u] < depth[v]) std::swap(u, v);
        u = kth_ancestor(u, depth[u] - depth[v]);
        if (u == v) return u;
        for (int k = _log - 1; k >= 0; k--) {
            if (up[k][u] != up[k][v]) {
                u = up[k][u];
                v = up[k][v];
            }
        }
        return parent[u];
    }

    int dist_edges(int u, int v) const {
        int w = lca(u, v);
        return depth[u] + depth[v] - 2 * depth[w];
    }

    T dist_cost(int u, int v) const {
        int w = lca(u, v);
        return dist[u] + dist[v] - dist[w] - dist[w];
    }

    int jump(int from, int to, int k) const {
        check_vertex(from);
        check_vertex(to);
        assert(0 <= k);
        int w = lca(from, to);
        int up_len = depth[from] - depth[w];
        int down_len = depth[to] - depth[w];
        if (up_len + down_len < k) return -1;
        if (k <= up_len) return kth_ancestor(from, k);
        return kth_ancestor(to, down_len - (k - up_len));
    }

    std::vector<int> path(int u, int v) const {
        check_vertex(u);
        check_vertex(v);
        int w = lca(u, v);
        std::vector<int> a, b;
        for (int x = u; x != w; x = parent[x]) a.push_back(x);
        a.push_back(w);
        for (int x = v; x != w; x = parent[x]) b.push_back(x);
        std::reverse(b.begin(), b.end());
        a.insert(a.end(), b.begin(), b.end());
        return a;
    }

    std::vector<int> path_edges(int u, int v) const {
        check_vertex(u);
        check_vertex(v);
        int w = lca(u, v);
        std::vector<int> a, b;
        for (int x = u; x != w; x = parent[x]) a.push_back(parent_edge[x]);
        for (int x = v; x != w; x = parent[x]) b.push_back(parent_edge[x]);
        std::reverse(b.begin(), b.end());
        a.insert(a.end(), b.begin(), b.end());
        return a;
    }

    std::pair<int, int> subtree_range(int v) const {
        check_vertex(v);
        return {tin[v], tout[v]};
    }

    std::vector<int> subtree_vertices(int v) const {
        check_vertex(v);
        return std::vector<int>(order.begin() + tin[v], order.begin() + tout[v]);
    }
};

}  // namespace tree
}  // namespace m1une


#line 9 "graph/tree/zero_one_on_tree.hpp"

namespace m1une {
namespace tree {

inline long long zero_one_on_tree(const std::vector<int>& parent,
                                  const std::vector<int>& value) {
    const int n = int(parent.size());
    assert(int(value.size()) == n);
    if (n == 0) return 0;

    int root = -1;
    std::vector<std::vector<int>> children(n);
    for (int v = 0; v < n; v++) {
        assert(value[v] == 0 || value[v] == 1);
        if (parent[v] == -1) {
            assert(root == -1);
            root = v;
        } else {
            assert(0 <= parent[v] && parent[v] < n && parent[v] != v);
            children[parent[v]].push_back(v);
        }
    }
    assert(root != -1);

    std::vector<int> stack(1, root);
    std::vector<char> visited(n, false);
    visited[root] = true;
    int visited_count = 0;
    while (!stack.empty()) {
        const int v = stack.back();
        stack.pop_back();
        visited_count++;
        for (int child : children[v]) {
            assert(!visited[child]);
            visited[child] = true;
            stack.push_back(child);
        }
    }
    assert(visited_count == n);

    struct Component {
        long long zeros;
        long long ones;
        int vertex;
    };
    struct Compare {
        bool operator()(const Component& lhs, const Component& rhs) const {
            const long long lhs_product = lhs.zeros * rhs.ones;
            const long long rhs_product = rhs.zeros * lhs.ones;
            if (lhs_product != rhs_product) return lhs_product < rhs_product;
            return lhs.vertex < rhs.vertex;
        }
    };

    std::vector<long long> zeros(n), ones(n);
    std::vector<int> dsu(n);
    std::set<Component, Compare> components;
    for (int v = 0; v < n; v++) {
        zeros[v] = value[v] == 0;
        ones[v] = value[v] == 1;
        dsu[v] = v;
        if (v != root) components.insert(Component{zeros[v], ones[v], v});
    }

    auto leader = [&](int v) {
        int result = v;
        while (dsu[result] != result) result = dsu[result];
        while (dsu[v] != v) {
            const int next = dsu[v];
            dsu[v] = result;
            v = next;
        }
        return result;
    };

    long long answer = 0;
    while (!components.empty()) {
        auto it = components.end();
        --it;
        const Component child = *it;
        components.erase(it);

        const int p = leader(parent[child.vertex]);
        if (p != root) {
            const int erased = int(components.erase(Component{zeros[p], ones[p], p}));
            assert(erased == 1);
        }

        answer += ones[p] * zeros[child.vertex];
        zeros[p] += zeros[child.vertex];
        ones[p] += ones[child.vertex];
        dsu[child.vertex] = p;

        if (p != root) components.insert(Component{zeros[p], ones[p], p});
    }
    return answer;
}

template <class T>
long long zero_one_on_tree(const m1une::graph::Graph<T>& graph,
                           const std::vector<int>& value, int root = 0) {
    const int n = graph.size();
    assert(int(value.size()) == n);
    if (n == 0) return 0;
    assert(0 <= root && root < n);
    assert(int(graph.edges().size()) == n - 1);

    RootedTree<T> rooted_tree(graph, root);
    assert(int(rooted_tree.order.size()) == n);
    return zero_one_on_tree(rooted_tree.parent, value);
}

}  // namespace tree
}  // namespace m1une

m1une's library

This documentation is automatically generated by online-judge-tools/verification-helper

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

Overview

Public Interface

Complexity

Example

Depends on

Required by

Verified with

Code

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

Overview

Public Interface

Complexity

Example

Depends on

Required by

Verified with

Code

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