#1013248 (C++23) No.2892 Lime and Karin

提出ソース

結果

問題	No.2892 Lime and Karin
ユーザー	hitonanode
提出日時	2024-09-25 00:31:44
言語	C++23 (gcc 12.3.0 + boost 1.83.0)
結果	AC
実行時間	434 ms / 8,000 ms
コード長	5,922 bytes
コンパイル時間	1,479 ms
コンパイル使用メモリ	117,792 KB
実行使用メモリ	31,260 KB
最終ジャッジ日時	2024-09-25 00:32:07
合計ジャッジ時間	14,025 ms
ジャッジサーバーID （参考情報）	judge4 / judge3

このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示

入力	結果	実行時間実行使用メモリ
testcase_00	AC	2 ms 6,812 KB
testcase_01	AC	2 ms 6,944 KB
testcase_02	AC	2 ms 6,940 KB
testcase_03	AC	2 ms 6,940 KB
testcase_04	AC	2 ms 6,944 KB
testcase_05	AC	2 ms 6,940 KB
testcase_06	AC	2 ms 6,940 KB
testcase_07	AC	2 ms 6,944 KB
testcase_08	AC	2 ms 6,944 KB
testcase_09	AC	2 ms 6,940 KB
testcase_10	AC	2 ms 6,944 KB
testcase_11	AC	2 ms 6,940 KB
testcase_12	AC	2 ms 6,940 KB
testcase_13	AC	2 ms 6,940 KB
testcase_14	AC	152 ms 10,624 KB
testcase_15	AC	141 ms 10,368 KB
testcase_16	AC	265 ms 14,392 KB
testcase_17	AC	49 ms 6,940 KB
testcase_18	AC	267 ms 14,440 KB
testcase_19	AC	434 ms 18,176 KB
testcase_20	AC	15 ms 6,944 KB
testcase_21	AC	180 ms 11,904 KB
testcase_22	AC	245 ms 14,232 KB
testcase_23	AC	87 ms 8,064 KB
testcase_24	AC	424 ms 18,316 KB
testcase_25	AC	400 ms 18,316 KB
testcase_26	AC	395 ms 18,436 KB
testcase_27	AC	419 ms 18,336 KB
testcase_28	AC	388 ms 18,312 KB
testcase_29	AC	387 ms 18,352 KB
testcase_30	AC	424 ms 18,312 KB
testcase_31	AC	387 ms 18,424 KB
testcase_32	AC	399 ms 18,440 KB
testcase_33	AC	378 ms 18,316 KB
testcase_34	AC	85 ms 16,944 KB
testcase_35	AC	88 ms 17,068 KB
testcase_36	AC	89 ms 17,068 KB
testcase_37	AC	90 ms 17,072 KB
testcase_38	AC	93 ms 16,832 KB
testcase_39	AC	385 ms 25,200 KB
testcase_40	AC	378 ms 30,664 KB
testcase_41	AC	399 ms 26,164 KB
testcase_42	AC	374 ms 31,260 KB
testcase_43	AC	377 ms 27,720 KB
testcase_44	AC	114 ms 9,088 KB
testcase_45	AC	337 ms 16,544 KB
testcase_46	AC	146 ms 10,112 KB
testcase_47	AC	197 ms 12,288 KB
testcase_48	AC	69 ms 7,424 KB
testcase_49	AC	240 ms 13,564 KB
testcase_50	AC	214 ms 17,940 KB
testcase_51	AC	215 ms 17,944 KB
testcase_52	AC	223 ms 17,816 KB
testcase_53	AC	220 ms 17,820 KB
testcase_54	AC	215 ms 17,816 KB

権限があれば一括ダウンロードができます

ソースコード

raw source code

#include <cassert>
#include <tuple>
#include <utility>
#include <vector>

struct CentroidDecomposition {
    int NO_PARENT = -1;
    int V;
    int E;
    std::vector<std::vector<std::pair<int, int>>> to; // (node_id, edge_id)
    std::vector<int> par;                             // parent node_id par[root] = -1
    std::vector<std::vector<int>> chi;                // children id's
    std::vector<int> subtree_size;                    // size of each subtree
    std::vector<int> available_edge;                  // If 0, ignore the corresponding edge.

    CentroidDecomposition(int v = 0) : V(v), E(0), to(v), par(v, NO_PARENT), chi(v), subtree_size(v) {}
    CentroidDecomposition(const std::vector<std::vector<int>> &to_) : CentroidDecomposition(to_.size()) {
        for (int i = 0; i < V; i++) {
            for (auto j : to_[i]) {
                if (i < j) { add_edge(i, j); }
            }
        }
    }

    void add_edge(int v1, int v2) {
        to[v1].emplace_back(v2, E), to[v2].emplace_back(v1, E), E++;
        available_edge.emplace_back(1);
    }

    int _dfs_fixroot(int now, int prv) {
        chi[now].clear(), subtree_size[now] = 1;
        for (auto nxt : to[now]) {
            if (nxt.first != prv and available_edge[nxt.second]) {
                par[nxt.first] = now, chi[now].push_back(nxt.first);
                subtree_size[now] += _dfs_fixroot(nxt.first, now);
            }
        }
        return subtree_size[now];
    }

    void fix_root(int root) {
        par[root] = NO_PARENT;
        _dfs_fixroot(root, -1);
    }

    //// Centroid Decpmposition ////
    std::vector<int> centroid_cand_tmp;
    void _dfs_detect_centroids(int now, int prv, int n) {
        bool is_centroid = true;
        for (auto nxt : to[now]) {
            if (nxt.first != prv and available_edge[nxt.second]) {
                _dfs_detect_centroids(nxt.first, now, n);
                if (subtree_size[nxt.first] > n / 2) is_centroid = false;
            }
        }
        if (n - subtree_size[now] > n / 2) is_centroid = false;
        if (is_centroid) centroid_cand_tmp.push_back(now);
    }
    std::pair<int, int> detect_centroids(int r) { // ([centroid_node_id1], ([centroid_node_id2]|-1))
        centroid_cand_tmp.clear();
        while (par[r] != NO_PARENT) r = par[r];
        int n = subtree_size[r];
        _dfs_detect_centroids(r, -1, n);
        if (centroid_cand_tmp.size() == 1)
            return std::make_pair(centroid_cand_tmp[0], -1);
        else
            return std::make_pair(centroid_cand_tmp[0], centroid_cand_tmp[1]);
    }

    std::vector<int> _cd_vertices;
    void _centroid_decomposition(int now) {
        fix_root(now);
        now = detect_centroids(now).first;
        _cd_vertices.emplace_back(now);
        /*
        do something
        */
        for (auto p : to[now]) {
            int nxt, eid;
            std::tie(nxt, eid) = p;
            if (available_edge[eid] == 0) continue;
            available_edge[eid] = 0;
            _centroid_decomposition(nxt);
        }
    }
    std::vector<int> centroid_decomposition(int x) {
        _cd_vertices.clear();
        _centroid_decomposition(x);
        return _cd_vertices;
    }
};

#include <algorithm>
#include <vector>

// 0-indexed BIT (binary indexed tree / Fenwick tree) (i : [0, len))
template <class T> struct BIT {
    int n;
    std::vector<T> data;
    BIT(int len = 0) : n(len), data(len) {}
    void reset() { std::fill(data.begin(), data.end(), T(0)); }
    void add(int pos, T v) { // a[pos] += v
        pos++;
        while (pos > 0 and pos <= n) data[pos - 1] += v, pos += pos & -pos;
    }
    T sum(int k) const { // a[0] + ... + a[k - 1]
        T res = 0;
        while (k > 0) res += data[k - 1], k -= k & -k;
        return res;
    }

    T sum(int l, int r) const { return sum(r) - sum(l); } // a[l] + ... + a[r - 1]

    template <class OStream> friend OStream &operator<<(OStream &os, const BIT &bit) {
        T prv = 0;
        os << '[';
        for (int i = 1; i <= bit.n; i++) {
            T now = bit.sum(i);
            os << now - prv << ',', prv = now;
        }
        return os << ']';
    }
};

#include <iostream>

int main() {
    int N;
    std::cin >> N;

    CentroidDecomposition cd(N);

    for (int e = 0; e < N - 1; ++e) {
        int u, v;
        std::cin >> u >> v;
        --u, --v;
        cd.add_edge(u, v);
    }

    std::vector<int> V(N);

    {
        std::string S;
        std::cin >> S;
        for (int i = 0; i < N; ++i) {
            V.at(i) = (S.at(i) - '0') * 2 - 1;
        }
    }

    std::vector<int> is_alive(N, 1);

    long long ret = 0;

    std::vector<int> cnt(N * 2 + 1);
    BIT<int> bit(N * 2 + 1);

    for (int c : cd.centroid_decomposition(0)) {
        is_alive.at(c) = 0;

        int lo = 0, hi = 0;

        auto addbit = [&](int w) -> void {
            lo = std::min(lo, w);
            hi = std::max(hi, w);
            cnt.at(N + w)++;
            bit.add(N + w, 1);
        };

        addbit(V.at(c));
        if (V.at(c) > 0) ret++;

        for (auto [nxt, _] : cd.to.at(c)) {
            if (!is_alive.at(nxt)) continue;

            std::vector<int> ws;

            auto dfs = [&](auto &&self, int now, int prv, int w) -> void {
                ws.push_back(w);
                // addbit(w);

                ret += bit.sum(-w + N + 1, N * 2 + 1);

                for (auto [nxt, _] : cd.to.at(now)) {
                    if (nxt == prv or !is_alive.at(nxt)) continue;
                    self(self, nxt, now, w + V.at(nxt));
                }
            };
            dfs(dfs, nxt, c, V.at(nxt));

            for (int w : ws) {
                w += V.at(c);
                addbit(w);
            }
        }

        for (int i = lo; i <= hi; ++i) {
            int j = i + N;
            bit.add(j, -cnt.at(j));
            cnt.at(j) = 0;
        }
    }

    std::cout << ret << '\n';
}

yukicoder

結果

テストケース

ソースコード