verify/yukicoder-3148.test.cpp

View this file on GitHub
Last update: 2026-05-31 17:59:37+09:00
Problem: https://yukicoder.me/problems/no/3148

Depends on

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

Code

// competitive-verifier: PROBLEM https://yukicoder.me/problems/no/3148

#include <iostream>
#include <string>
#include <utility>
#include <vector>

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

int main() {
    int n;
    std::string s;
    std::cin >> n >> s;

    const int root = n;
    std::vector<long long> c0(n + 1, 0), c1(n + 1, 0);
    long long sum = 0;
    for (int i = 0; i < n; ++i) {
        std::cin >> c0[i];
        c1[i] = 1;
        sum += c0[i];
    }

    std::vector<std::pair<int, int>> edges;
    edges.reserve(n);
    std::vector<int> stack;
    stack.reserve(n);
    int id = 0;
    for (char ch : s) {
        if (ch == '(') {
            const int v = id++;
            const int parent = stack.empty() ? root : stack.back();
            edges.push_back({parent, v});
            stack.push_back(v);
        } else {
            stack.pop_back();
        }
    }

    const long long answer =
        solve_01_on_tree<long long>(n + 1, edges, c0, c1, root) + sum;
    std::cout << answer << '\n';
    return 0;
}

#line 1 "verify/yukicoder-3148.test.cpp"
// competitive-verifier: PROBLEM https://yukicoder.me/problems/no/3148

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

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


#line 9 "verify/yukicoder-3148.test.cpp"

int main() {
    int n;
    std::string s;
    std::cin >> n >> s;

    const int root = n;
    std::vector<long long> c0(n + 1, 0), c1(n + 1, 0);
    long long sum = 0;
    for (int i = 0; i < n; ++i) {
        std::cin >> c0[i];
        c1[i] = 1;
        sum += c0[i];
    }

    std::vector<std::pair<int, int>> edges;
    edges.reserve(n);
    std::vector<int> stack;
    stack.reserve(n);
    int id = 0;
    for (char ch : s) {
        if (ch == '(') {
            const int v = id++;
            const int parent = stack.empty() ? root : stack.back();
            edges.push_back({parent, v});
            stack.push_back(v);
        } else {
            stack.pop_back();
        }
    }

    const long long answer =
        solve_01_on_tree<long long>(n + 1, edges, c0, c1, root) + sum;
    std::cout << answer << '\n';
    return 0;
}

Test cases

Env	Name	Status	Elapsed	Memory
g++	1_sample01	AC	2 ms	4 MB
g++	1_sample02	AC	2 ms	4 MB
g++	1_sample03	AC	2 ms	4 MB
g++	2_small01	AC	2 ms	4 MB
g++	2_small02	AC	2 ms	4 MB
g++	2_small03	AC	2 ms	3 MB
g++	2_small04	AC	2 ms	4 MB
g++	2_small05	AC	2 ms	3 MB
g++	3_random01	AC	62 ms	7 MB
g++	3_random02	AC	111 ms	10 MB
g++	3_random03	AC	16 ms	4 MB
g++	3_random04	AC	109 ms	10 MB
g++	3_random05	AC	100 ms	9 MB
g++	3_random06	AC	165 ms	12 MB
g++	3_random07	AC	216 ms	16 MB
g++	3_random08	AC	250 ms	18 MB
g++	3_random09	AC	126 ms	11 MB
g++	3_random10	AC	212 ms	17 MB
g++	3_random11	AC	197 ms	16 MB
g++	3_random12	AC	94 ms	9 MB
g++	3_random13	AC	33 ms	5 MB
g++	3_random14	AC	77 ms	8 MB
g++	3_random15	AC	50 ms	6 MB
g++	3_random16	AC	161 ms	12 MB
g++	3_random17	AC	67 ms	7 MB
g++	3_random18	AC	172 ms	14 MB
g++	3_random19	AC	73 ms	7 MB
g++	3_random20	AC	197 ms	15 MB
g++	4_max01	AC	264 ms	19 MB
g++	4_max02	AC	266 ms	19 MB
g++	4_max03	AC	295 ms	18 MB
g++	4_max04	AC	334 ms	19 MB
g++	4_max05	AC	141 ms	11 MB
g++	4_max06	AC	153 ms	11 MB
clang++	1_sample01	AC	2 ms	4 MB
clang++	1_sample02	AC	2 ms	4 MB
clang++	1_sample03	AC	2 ms	4 MB
clang++	2_small01	AC	2 ms	3 MB
clang++	2_small02	AC	2 ms	3 MB
clang++	2_small03	AC	2 ms	4 MB
clang++	2_small04	AC	2 ms	3 MB
clang++	2_small05	AC	2 ms	4 MB
clang++	3_random01	AC	55 ms	7 MB
clang++	3_random02	AC	99 ms	10 MB
clang++	3_random03	AC	15 ms	4 MB
clang++	3_random04	AC	97 ms	10 MB
clang++	3_random05	AC	89 ms	9 MB
clang++	3_random06	AC	147 ms	12 MB
clang++	3_random07	AC	192 ms	16 MB
clang++	3_random08	AC	224 ms	17 MB
clang++	3_random09	AC	113 ms	11 MB
clang++	3_random10	AC	193 ms	16 MB
clang++	3_random11	AC	174 ms	15 MB
clang++	3_random12	AC	83 ms	9 MB
clang++	3_random13	AC	29 ms	5 MB
clang++	3_random14	AC	68 ms	8 MB
clang++	3_random15	AC	44 ms	7 MB
clang++	3_random16	AC	144 ms	12 MB
clang++	3_random17	AC	59 ms	7 MB
clang++	3_random18	AC	152 ms	14 MB
clang++	3_random19	AC	64 ms	7 MB
clang++	3_random20	AC	178 ms	15 MB
clang++	4_max01	AC	242 ms	18 MB
clang++	4_max02	AC	236 ms	18 MB
clang++	4_max03	AC	253 ms	18 MB
clang++	4_max04	AC	291 ms	18 MB
clang++	4_max05	AC	123 ms	11 MB
clang++	4_max06	AC	137 ms	11 MB