結果

問題 No.2892 Lime and Karin
ユーザー fuppy_kyoprofuppy_kyopro
提出日時 2024-09-13 23:05:01
言語 C++17
(gcc 12.3.0 + boost 1.83.0)
結果
WA  
実行時間 -
コード長 18,545 bytes
コンパイル時間 3,569 ms
コンパイル使用メモリ 244,564 KB
実行使用メモリ 55,312 KB
最終ジャッジ日時 2024-09-13 23:05:13
合計ジャッジ時間 11,402 ms
ジャッジサーバーID
(参考情報)
judge6 / judge1
このコードへのチャレンジ
(要ログイン)

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 2 ms
6,816 KB
testcase_01 AC 2 ms
6,816 KB
testcase_02 AC 2 ms
6,944 KB
testcase_03 AC 2 ms
6,944 KB
testcase_04 WA -
testcase_05 WA -
testcase_06 WA -
testcase_07 WA -
testcase_08 WA -
testcase_09 WA -
testcase_10 WA -
testcase_11 WA -
testcase_12 WA -
testcase_13 WA -
testcase_14 WA -
testcase_15 WA -
testcase_16 WA -
testcase_17 WA -
testcase_18 WA -
testcase_19 WA -
testcase_20 WA -
testcase_21 WA -
testcase_22 WA -
testcase_23 WA -
testcase_24 WA -
testcase_25 WA -
testcase_26 WA -
testcase_27 WA -
testcase_28 WA -
testcase_29 WA -
testcase_30 WA -
testcase_31 WA -
testcase_32 WA -
testcase_33 WA -
testcase_34 AC 123 ms
55,312 KB
testcase_35 WA -
testcase_36 AC 129 ms
55,312 KB
testcase_37 AC 132 ms
55,312 KB
testcase_38 AC 129 ms
55,312 KB
testcase_39 AC 183 ms
48,788 KB
testcase_40 AC 190 ms
49,520 KB
testcase_41 WA -
testcase_42 WA -
testcase_43 WA -
testcase_44 AC 66 ms
20,240 KB
testcase_45 AC 168 ms
41,572 KB
testcase_46 AC 78 ms
22,884 KB
testcase_47 AC 113 ms
29,164 KB
testcase_48 AC 46 ms
14,944 KB
testcase_49 AC 133 ms
32,976 KB
testcase_50 AC 152 ms
50,904 KB
testcase_51 AC 156 ms
51,032 KB
testcase_52 WA -
testcase_53 WA -
testcase_54 WA -
権限があれば一括ダウンロードができます

ソースコード

diff #

/*
#pragma GCC target("avx2")
#pragma GCC optimize("O3")
#pragma GCC optimize("unroll-loops")
//*/

// #include <atcoder/all>
#include <bits/stdc++.h>

using namespace std;
// using namespace atcoder;

// #define _GLIBCXX_DEBUG

#define DEBUG(x) cerr << #x << ": " << x << endl;
#define DEBUG_VEC(v)                                        \
    cerr << #v << ":";                                      \
    for (int iiiiiiii = 0; iiiiiiii < v.size(); iiiiiiii++) \
        cerr << " " << v[iiiiiiii];                         \
    cerr << endl;
#define DEBUG_MAT(v)                                \
    cerr << #v << endl;                             \
    for (int iv = 0; iv < v.size(); iv++) {         \
        for (int jv = 0; jv < v[iv].size(); jv++) { \
            cerr << v[iv][jv] << " ";               \
        }                                           \
        cerr << endl;                               \
    }
typedef long long ll;
// #define int ll

#define vi vector<int>
#define vl vector<ll>
#define vii vector<vector<int>>
#define vll vector<vector<ll>>
#define pii pair<int, int>
#define pis pair<int, string>
#define psi pair<string, int>
#define pll pair<ll, ll>
template <class S, class T>
pair<S, T> operator+(const pair<S, T> &s, const pair<S, T> &t) {
    return pair<S, T>(s.first + t.first, s.second + t.second);
}
template <class S, class T>
pair<S, T> operator-(const pair<S, T> &s, const pair<S, T> &t) { return pair<S, T>(s.first - t.first, s.second - t.second); }
template <class S, class T>
ostream &operator<<(ostream &os, pair<S, T> p) {
    os << "(" << p.first << ", " << p.second << ")";
    return os;
}
#define rep(i, n) for (int i = 0; i < (int)(n); i++)
#define rep1(i, n) for (int i = 1; i <= (int)(n); i++)
#define rrep(i, n) for (int i = (int)(n) - 1; i >= 0; i--)
#define rrep1(i, n) for (int i = (int)(n); i > 0; i--)
#define REP(i, a, b) for (int i = a; i < b; i++)
#define in(x, a, b) (a <= x && x < b)
#define all(c) c.begin(), c.end()
void YES(bool t = true) {
    cout << (t ? "YES" : "NO") << endl;
}
void Yes(bool t = true) { cout << (t ? "Yes" : "No") << endl; }
void yes(bool t = true) { cout << (t ? "yes" : "no") << endl; }
void NO(bool t = true) { cout << (t ? "NO" : "YES") << endl; }
void No(bool t = true) { cout << (t ? "No" : "Yes") << endl; }
void no(bool t = true) { cout << (t ? "no" : "yes") << endl; }
template <class T>
bool chmax(T &a, const T &b) {
    if (a < b) {
        a = b;
        return 1;
    }
    return 0;
}
template <class T>
bool chmin(T &a, const T &b) {
    if (a > b) {
        a = b;
        return 1;
    }
    return 0;
}
template <class T, class U>
T ceil_div(T a, U b) {
    return (a + b - 1) / b;
}
template <class T>
T safe_mul(T a, T b) {
    if (a == 0 || b == 0) return 0;
    if (numeric_limits<T>::max() / a < b) return numeric_limits<T>::max();
    return a * b;
}
#define UNIQUE(v) v.erase(std::unique(v.begin(), v.end()), v.end());
const ll inf = 1000000001;
const ll INF = (ll)1e18 + 1;
const long double pi = 3.1415926535897932384626433832795028841971L;
int popcount(ll t) { return __builtin_popcountll(t); }
vector<int> gen_perm(int n) {
    vector<int> ret(n);
    iota(all(ret), 0);
    return ret;
}
// int dx[4] = {1, 0, -1, 0}, dy[4] = {0, 1, 0, -1};
// int dx2[8] = { 1,1,0,-1,-1,-1,0,1 }, dy2[8] = { 0,1,1,1,0,-1,-1,-1 };
vi dx = {0, 0, -1, 1}, dy = {-1, 1, 0, 0};
vi dx2 = {1, 1, 0, -1, -1, -1, 0, 1}, dy2 = {0, 1, 1, 1, 0, -1, -1, -1};
struct Setup_io {
    Setup_io() {
        ios_base::sync_with_stdio(0), cin.tie(0), cout.tie(0);
        cout << fixed << setprecision(25);
        cerr << fixed << setprecision(25);
    }
} setup_io;
// constexpr ll MOD = 1000000007;
constexpr ll MOD = 998244353;
// #define mp make_pair

enum Type { Vertex,
            Compress,
            Rake,
            AddEdge,
            AddVertex };

struct StaticTopTree {
    // https://atcoder.jp/contests/abc351/editorial/9868
    // https://atcoder.jp/contests/abc351/submissions/52777033 を元に、自分好みに書き方を変えたりコメントを追加したり辺クエリに対応したりした
    // 全方位木 DP のライブラリと考え方が似ているのでそちらも参照すると良い

    // G は n 頂点の木を表す連結リストであるとする。
    // 木を HLD のように heavy path と light path に分解する。この時、親からの辺が light edge である頂点と heavy edge である頂点の二種類に分類できる。
    // 前者の頂点を軽頂点、後者の頂点を重頂点と呼ぶ。ただし、根は重頂点であるとする。
    // 頂点も辺も存在しない状態から、以下の操作を行うことで元の木を構築することを考える。この際、頂点が各操作を表すような操作列を表す木も同時に構築する。
    //
    // A-1. 子に軽頂点を持たない頂点を追加する(Vertex)
    // B. 軽頂点が根のクラスタ(部分木的なもの)に親辺を追加する。親頂点は追加しないことに注意(AddEdge)
    // C. 複数の「B. で作られたもの」をマージする(Rake)
    // A-2. 「C. で作られたもの」の親頂点を追加し、「C. で作られたもの」とくっつける(AddVertex)
    // D. 重頂点とその親頂点をマージすることで heavy path を伸ばす(Compress)
    //
    // 実装上は C, D は二つのクラスタのマージを繰り返すことで表現するため、「操作列を表す木」は二分木となる。
    // この際、マージする順番を工夫することで完全二分木に近づけることができ、木の深さを O(log n) に抑えることができる。

    vector<vector<pair<int, int>>> G; // 元のグラフ、(to, edge_idx) のリストで表され、親への辺は含まれない
    int root;                         // 元のグラフでの根
    vector<int> par_edge;             // 元のグラフで各頂点の親への辺のインデックス
    int stt_root;                     // 操作列を表す木の根
    vector<int> P, L, R;              // 操作列を表す木における、親、左の子、右の子
    vector<Type> T;                   // 操作列を表す木における、各頂点の操作タイプ
    vector<int> merged_edge;          // compress / add_edge の操作を表す頂点において、その辺のインデックス
    int add_buf;                      // 操作列を表す木における、新たに追加する頂点のインデックス

    StaticTopTree(const vector<vector<pair<int, int>>> &_G = {}, int _root = 0) : G(_G), root(_root) {
        int n = G.size();
        if (n == 0) {
            return;
        }
        par_edge.resize(n, -1);

        P.resize(4 * n, -1), L.resize(4 * n, -1), R.resize(4 * n, -1);
        T.resize(4 * n, Type::Vertex);
        merged_edge.resize(4 * n, -1);
        add_buf = n;
        build();
    }

    int size() const {
        // 操作列を表す木のサイズを返す
        return add_buf;
    }

  private:
    int dfs(int c, int par = -1) {
        // 部分木のサイズを計算し、最も大きい部分木をリストの先頭に持ってくる
        // また、各頂点の親への辺のインデックスを記録する
        int s = 1, best = 0;
        for (int i = 0; i < (int)G[c].size(); i++) {
            auto [d, ei] = G[c][i];
            par_edge[d] = ei;
            int t = dfs(d, c);
            s += t;
            if (best < t) best = t, swap(G[c][i], G[c][0]);
        }
        return s;
    }
    int add(int k, int l, int r, Type t) {
        if (k == -1) k = add_buf++;
        P[k] = -1, L[k] = l, R[k] = r, T[k] = t;
        if (l != -1) P[l] = k;
        if (r != -1) P[r] = k;
        return k;
    }

    // 以下の pair<int, int> を返す関数は、(操作に対応する頂点のインデックス, 部分木に含まれる実頂点の数) を返す
    // ただし、操作に対応する頂点が作られない場合は (-1, 0) を返す
    pair<int, int> merge(const vector<pair<int, int>> &a, Type t) {
        assert(a.size() > 0);
        if (a.size() == 1) return a[0];
        int u = 0;
        for (auto &[_, s] : a)
            u += s;

        // 二つに振り分けたあとにマージする
        vector<pair<int, int>> b, c;
        for (auto &[i, s] : a)
            (u > s ? b : c).emplace_back(i, s), u -= s * 2;

        auto [i, si] = merge(b, t);
        auto [j, sj] = merge(c, t);

        auto ret = pair<int, int>(add(-1, i, j, t), si + sj);
        if (t == Type::Compress) {
            // compress の操作を表す頂点において、その辺のインデックスを記録する
            // c[0] が今回マージされる子頂点であり、c[0] が Vertex, AddVertex であることは `compress` 関数の仕様から保証される
            merged_edge[ret.first] = par_edge[c[0].first];
        }
        return ret;
    }
    pair<int, int> compress(int i) {
        vector<pair<int, int>> chs;
        while (true) {
            chs.push_back(add_vertex(i));
            if (G[i].size() == 0) {
                break;
            }
            i = G[i][0].first;
        }
        return merge(chs, Type::Compress);
    }
    pair<int, int> rake(int i) {
        vector<pair<int, int>> chs;
        for (int j = 1; j < (int)G[i].size(); j++)
            chs.push_back(add_edge(G[i][j].first));
        return chs.empty() ? make_pair(-1, 0) : merge(chs, Type::Rake);
    }
    pair<int, int> add_edge(int i) {
        // i の親は light path であることが保証される
        auto [j, sj] = compress(i);
        auto ret = pair<int, int>(add(-1, j, -1, Type::AddEdge), sj);
        // add_edge の操作を表す頂点において、その辺のインデックスを記録する
        // i が Vertex, AddVertex であることは `rake` 関数の仕様から保証される
        merged_edge[ret.first] = par_edge[i];
        return ret;
    }
    pair<int, int> add_vertex(int i) {
        auto [j, sj] = rake(i);
        return {add(i, j, -1, j == -1 ? Type::Vertex : Type::AddVertex), sj + 1};
    }
    void build() {
        dfs(root);
        auto [i, n] = compress(root);
        stt_root = i;
    }
};
template <class E, class V>
struct StaticTopTreeDp {
    // static top tree を用いて以下を実現する
    // - 各頂点や辺がある値を持つ木に関する木 DP を行う
    // - 各頂点や辺が持つ値が変わった時に O(log n) で木 DP の値を更新する
    //   - https://atcoder.jp/contests/abc351/submissions/52951833
    //   - https://judge.yosupo.jp/submission/205112
    // - 部分木のサイズ程度の長さの列を値に持つ DP は、分割統治的な考え方で O(n (log n)^2) に高速化できる
    //   - https://atcoder.jp/contests/abc269/submissions/52956082
    //
    // 全方位木 DP のライブラリと考え方が似ているのでそちらも参照すると良い
    // E: クラスタに親辺を追加したもの、それらをマージしたものが持つべき答えの型。AddEdge, Rake に対応。
    // V: 頂点を根に持つクラスタが持つべき答えの型。Vertex, AddVertex, Compress に対応。
    //    ここで、AddVertex と Compress に対応するため、重頂点の子の情報は持たないことに注意。

    StaticTopTree stt;

    function<V(int)> vertex;         // 頂点のインデックス
    function<E(V, int)> add_edge;    // (クラスタの答え, 辺のインデックス)
    function<E(E, E)> rake;          // (マージ対象の左クラスタの答え, 右)
    function<V(E, int)> add_vertex;  // (クラスタの答え, 頂点のインデックス)
    function<V(V, V, int)> compress; // (マージ対象の親クラスタの答え, 子, 辺のインデックス)

    vector<int> child; // 辺ごとに、子の方の頂点を記録する

    vector<E> dp_e; // stt の各頂点に対応するクラスタの答え。頂点が AddEdge, Rake でない場合は使用されない想定。
    vector<V> dp_v; // stt の各頂点に対応するクラスタの答え。頂点が Vertex, AddVertex, Compress でない場合は使用されない想定。

    StaticTopTreeDp(vector<vector<pair<int, int>>> G,
                    function<V(int)> vertex,
                    function<E(V, int)> add_edge,
                    function<E(E, E)> rake,
                    function<V(E, int)> add_vertex,
                    function<V(V, V, int)> compress,
                    int root = 0)
        : vertex(vertex), add_edge(add_edge), rake(rake), add_vertex(add_vertex), compress(compress) {

        child.resize(G.size() - 1, -1);
        preprocess_graph(G, root, -1);

        stt = StaticTopTree(G, root);

        dp_e.resize(stt.size());
        dp_v.resize(stt.size());
    }

    V fill_dp_all() {
        // DP を全て埋める。最初に一度だけ呼ばれる想定。
        dfs(stt.stt_root);
        return dp_v[stt.stt_root];
    }

    V update_vertex(int u) {
        // 頂点 u の値が変わった後に、その影響を木 DP に反映する。O(log n)
        while (u != -1) {
            update(u);
            u = stt.P[u];
        }
        return dp_v[stt.stt_root];
    }

    V update_edge(int e) {
        // 辺 e の値が変わった後に、その影響を木 DP に反映する。O(log n)
        return update_vertex(child[e]);
    }

  private:
    void preprocess_graph(vector<vector<pair<int, int>>> &G, int now, int par) {
        int par_idx = -1;
        for (int i = 0; i < (int)G[now].size(); i++) {
            if (G[now][i].first == par) {
                par_idx = i;
            } else {
                child[G[now][i].second] = G[now][i].first;
                preprocess_graph(G, G[now][i].first, now);
            }
        }
        if (par_idx != -1) {
            swap(G[now][par_idx], G[now].back());
            G[now].pop_back();
        }
    }

    void update(int u) {
        if (stt.T[u] == Type::Vertex) {
            dp_v[u] = vertex(u);
        } else if (stt.T[u] == Type::Compress) {
            dp_v[u] = compress(dp_v[stt.L[u]], dp_v[stt.R[u]], stt.merged_edge[u]);
        } else if (stt.T[u] == Type::Rake) {
            dp_e[u] = rake(dp_e[stt.L[u]], dp_e[stt.R[u]]);
        } else if (stt.T[u] == Type::AddEdge) {
            dp_e[u] = add_edge(dp_v[stt.L[u]], stt.merged_edge[u]);
        } else {
            dp_v[u] = add_vertex(dp_e[stt.L[u]], u);
        }
    }

    void dfs(int now) {
        if (stt.L[now] != -1) dfs(stt.L[now]);
        if (stt.R[now] != -1) dfs(stt.R[now]);
        update(now);
    }
};

vector<vector<pii>> G;
vector<pii> es;
vi a;

struct V {
    int root;
    vector<int> num;
    int offset;

    V() {}
    V(int root, vector<int> num, int offset) : root(root), num(num), offset(offset) {}
};

struct E {
    int root;
    vector<int> num;
    int offset;

    E() {}
    E(int root, vector<int> num, int offset) : root(root), num(num), offset(offset) {}
};

vi add(vi a, vi b) {
    if (a.size() < b.size()) {
        swap(a, b);
    }
    for (int i = 0; i < b.size(); i++) {
        a[i] += b[i];
    }
    return a;
}

ll ans = 0;
V vertex(int i) {
    if (a[i] == 1) {
        ans++;
    }
    return V(i, {1}, a[i]);
}
E add_edge(V v, int ei) {
    auto [a, b] = es[ei];
    assert(v.root == a || v.root == b);
    int r;
    if (v.root == a) {
        r = b;
    } else {
        r = a;
    }

    return E(r, v.num, v.offset);
}
E rake(E l, E r) {
    assert(l.root == r.root);

    // DEBUG(l.root);
    // DEBUG_VEC(l.num);
    // DEBUG(l.offset);

    // DEBUG(r.root);
    // DEBUG_VEC(r.num);
    // DEBUG(r.offset);

    ll add = 0;
    int def = a[l.root];
    rep(i, l.num.size()) {
        rep(j, r.num.size()) {
            if (i + l.offset + j + r.offset + def > 0) {
                add += l.num[i] * r.num[j];
            }
        }
    }
    // DEBUG(add);
    ans += add;

    if (l.offset > r.offset) {
        swap(l, r);
    }
    for (int i = 0; i < r.num.size(); i++) {
        int ori_idx = i + r.offset;
        int new_idx = ori_idx - l.offset;

        while (l.num.size() <= new_idx) {
            l.num.push_back(0);
        }
        l.num[new_idx] += r.num[i];
    }

    return E(l.root, l.num, l.offset);
}
V add_vertex(E d, int i) {
    assert(d.root == i);
    if (a[i] == 1) {
        ans++;
    }

    // DEBUG(d.root);
    // DEBUG_VEC(d.num);
    // DEBUG(d.offset);

    ll add = 0;
    rep(j, d.num.size()) {
        if (j + d.offset + a[i] > 0) {
            add += d.num[j];
        }
    }
    // DEBUG(add);
    ans += add;

    vi num1 = d.num, num2 = {1};
    int offset1 = d.offset + a[i], offset2 = a[i];

    if (offset1 > offset2) {
        swap(num1, num2);
        swap(offset1, offset2);
    }

    for (int i = 0; i < num2.size(); i++) {
        int ori_idx = i + offset2;
        int new_idx = ori_idx - offset1;

        while (num1.size() <= new_idx) {
            num1.push_back(0);
        }
        num1[new_idx] += num2[i];
    }

    return V(i, num1, offset1);
};
V compress(V p, V c, int ei) {
    // DEBUG(p.root);
    // DEBUG_VEC(p.num);
    // DEBUG(p.offset);

    // DEBUG(c.root);
    // DEBUG_VEC(c.num);
    // DEBUG(c.offset);

    ll add = 0;
    rep(i, p.num.size()) {
        rep(j, c.num.size()) {
            if (i + p.offset + j + c.offset > 0) {
                add += p.num[i] * c.num[j];
            }
        }
    }
    // DEBUG(add);
    ans += add;

    c.offset += a[p.root];

    int root = p.root;

    if (p.offset > c.offset) {
        swap(p, c);
    }
    for (int i = 0; i < c.num.size(); i++) {
        int ori_idx = i + c.offset;
        int new_idx = ori_idx - p.offset;

        while (p.num.size() <= new_idx) {
            p.num.push_back(0);
        }
        p.num[new_idx] += c.num[i];
    }

    return V(root, p.num, p.offset);
}

signed main() {
    int n;
    cin >> n;
    G.resize(n);
    rep(i, n - 1) {
        int a, b;
        cin >> a >> b;
        a--, b--;
        G[a].push_back({b, i});
        G[b].push_back({a, i});
        es.emplace_back(a, b);
    }

    string s;
    cin >> s;
    rep(i, n) {
        if (s[i] == '0') {
            a.push_back(-1);
        } else {
            a.push_back(1);
        }
    }

    StaticTopTreeDp<E, V> sttd(G, vertex, add_edge, rake, add_vertex, compress);
    sttd.fill_dp_all();

    cout << ans << endl;
}
0