結果

問題 No.399 動的な領主
ユーザー MazesobaMazesoba
提出日時 2023-03-21 19:56:01
言語 C++17
(gcc 12.3.0 + boost 1.83.0)
結果
AC  
実行時間 646 ms / 2,000 ms
コード長 8,381 bytes
コンパイル時間 5,973 ms
コンパイル使用メモリ 275,708 KB
実行使用メモリ 28,012 KB
最終ジャッジ日時 2023-10-18 18:33:46
合計ジャッジ時間 12,105 ms
ジャッジサーバーID
（参考情報） 		judge13 / judge12
このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 2 ms
4,348 KB
testcase_01 AC 2 ms
4,348 KB
testcase_02 AC 2 ms
4,348 KB
testcase_03 AC 2 ms
4,348 KB
testcase_04 AC 4 ms
4,348 KB
testcase_05 AC 36 ms
5,112 KB
testcase_06 AC 637 ms
19,300 KB
testcase_07 AC 624 ms
19,300 KB
testcase_08 AC 627 ms
19,300 KB
testcase_09 AC 628 ms
19,300 KB
testcase_10 AC 5 ms
4,348 KB
testcase_11 AC 29 ms
5,112 KB
testcase_12 AC 450 ms
19,564 KB
testcase_13 AC 433 ms
19,564 KB
testcase_14 AC 106 ms
28,012 KB
testcase_15 AC 146 ms
28,012 KB
testcase_16 AC 233 ms
24,052 KB
testcase_17 AC 646 ms
19,300 KB
testcase_18 AC 622 ms
19,300 KB
権限があれば一括ダウンロードができます

ソースコード

diff # 

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

using namespace std;

// #define DISABLE_PRINT

#if defined(ENABLE_PRINT) && !defined(DISABLE_PRINT)

#define P(...) fprintf(stderr, __VA_ARGS__)
#define LP fprintf(stderr, "L: %d\n", __LINE__)

#else

#define P(...) ((void)0)
#define LP ((void)0)

#endif

#define rep(i, n) for (int i = 0; i < (int)(n); ++i)
#define ALL(x) x.begin(), x.end()

using ll = long long;
using ull = unsigned long long;

const int INF = 1100100100;
const ll INFLL = 1100100100100100100;

class HLD {
private:
    int _N;
    int _edgeCount{};
    struct Edge {
        int to;
        int idx;
    };
    vector<vector<Edge>> _g;

    vector<int> _hld;
    vector<int> _hldEdge;
    vector<int> _viToHi;
    vector<int> _eiToHi;
    vector<int> _leaders;
    vector<int> _depths;
    vector<int> _sizes;
    vector<int> _parents;

    int CalcSizeAndDepth(int depth, int x, int p) {
        int ans = 1;
        _depths[x] = depth;
        _parents[x] = p;
        for (auto nx : _g[x]) {
            if (nx.to == p) continue;
            ans += CalcSizeAndDepth(depth + 1, nx.to, x);
        }
        _sizes[x] = ans;
        return ans;
    }

    void SetHld(int x, int p, int baseIdx, int leader) {
        _hld[baseIdx] = x;
        _leaders[x] = leader;
        _viToHi[x] = baseIdx;

        int hi = -1;
        int ls = 0;
        rep(i, (int)_g[x].size()) {
            auto ci = _g[x][i];
            if (ci.to == p) {
                _hldEdge[baseIdx] = ci.idx;
                _eiToHi[ci.idx] = baseIdx;
                continue;
            }
            auto cs = _sizes[ci.to];
            if (cs > ls) {
                hi = i;
                ls = cs;
            }
        }
        if (hi == -1) return;

        baseIdx++;

        // 長い子をつなげる
        SetHld(_g[x][hi].to, x, baseIdx, leader);
        baseIdx += _sizes[_g[x][hi].to];

        // 短い子の処理
        rep(i, (int)_g[x].size()) {
            if (i == hi) continue;
            auto ci = _g[x][i];
            if (ci.to == p) continue;
            SetHld(ci.to, x, baseIdx, ci.to);
            baseIdx += _sizes[ci.to];
        }
    }

public:
    HLD(int N)
        : _N(N), _g(N), _hld(N), _hldEdge(N, -1), _viToHi(N), _eiToHi(N, -1), _leaders(N), _depths(N), _sizes(N),
          _parents(N) {
    }

    const vector<int> &Hld() const {
        return _hld;
    }

    const vector<int> &ViToHi() const {
        return _viToHi;
    }

    const vector<int> &EiToHi() const {
        return _eiToHi;
    }

    const vector<int> &HldEdge() const {
        return _hldEdge;
    }

    void Build() {
        assert(_edgeCount == _N - 1);

        CalcSizeAndDepth(0, 0, -1);

        P("Sizes: ");
        rep(i, _N) {
            P("%2d ", _sizes[i]);
        }
        P("\n");
        P("Depth: ");
        rep(i, _N) {
            P("%2d ", _depths[i]);
        }
        P("\n");

        SetHld(0, -1, 0, 0);

        P("Hld  : ");
        rep(i, _N) {
            P("%2d ", _hld[i]);
        }
        P("\n");

        P("HldE : ");
        rep(i, _N) {
            P("%2d ", _hldEdge[i]);
        }
        P("\n");

        P("Lead : ");
        rep(i, _N) {
            P("%2d ", _leaders[i]);
        }
        P("\n");

        P("vToH : ");
        rep(i, _N) {
            P("%2d ", _viToHi[i]);
        }
        P("\n");

        P("eToH : ");
        rep(i, _N) {
            P("%2d ", _eiToHi[i]);
        }
        P("\n");
    }

    void AddEdge(int u, int v) {
        assert(_edgeCount < _N - 1);
        assert(u >= 0 && u < _N);
        assert(v >= 0 && v < _N);
        _g[u].push_back({v, _edgeCount});
        _g[v].push_back({u, _edgeCount});
        _edgeCount++;
    }

    int Lca(int u, int v) const {
        assert(u >= 0 && u < _N);
        assert(v >= 0 && v < _N);
        using PT = pair<int, int>; // depth, vi;
        LP;
        PT l = {_depths[_leaders[u]], u};
        PT r = {_depths[_leaders[v]], v};
        LP;
        while (l != r) {
            P("l: %2d, %2d\n", l.first, l.second);
            P("r: %2d, %2d\n", r.first, r.second);
            P("----\n");
            if (l > r) swap(l, r);
            if (_leaders[l.second] == _leaders[r.second]) {
                r = l;
                continue;
            }
            auto to = _parents[_leaders[r.second]];
            r = {_depths[_leaders[to]], to};
        }
        return l.second;
    }

    template <typename T, typename Q, typename F>
    T Query(int u, int v, Q query, F merge, bool edge = false) const {
        assert(u >= 0 && u < _N);
        assert(v >= 0 && v < _N);

        struct E {
            T state;
            int depth;
            int vi;
            bool operator>(const E &r) const {
                return depth > r.depth;
            }
            void Dump() const {
                P("d: %d, v: %d\n", depth, vi);
            }
        };

        E l = {T(), _depths[_leaders[u]], u};
        E r = {T(), _depths[_leaders[v]], v};
        while (true) {
            if (l > r) swap(l, r);
            if (_leaders[l.vi] == _leaders[r.vi]) {
                auto li = _viToHi[l.vi];
                auto ri = _viToHi[r.vi];
                if (li > ri) swap(li, ri);
                auto add = edge ? 1 : 0;
                auto ta = query(li + add, ri + 1, r.state);
                auto ans = merge(ta, l.state);
                return ans;
            }
            auto li = _viToHi[_leaders[r.vi]];
            auto ri = _viToHi[r.vi];
            auto t = query(li, ri + 1, r.state);
            auto to = _parents[_leaders[r.vi]];
            r = {t, _depths[_leaders[to]], to};
        }
    }

    template <typename F>
    void Update(int u, int v, F update) const {
        struct E {
            int depth;
            int vi;
            bool operator>(const E &r) const {
                return depth > r.depth;
            }
            void Dump() const {
                P("d: %d, v: %d\n", depth, vi);
            }
        };

        E l = {_depths[_leaders[u]], u};
        E r = {_depths[_leaders[v]], v};
        while (true) {
            if (l > r) swap(l, r);
            if (_leaders[l.vi] == _leaders[r.vi]) {
                auto li = _viToHi[l.vi];
                auto ri = _viToHi[r.vi];
                if (li > ri) swap(li, ri);
                update(li, ri + 1);
                return;
            }
            auto li = _viToHi[_leaders[r.vi]];
            auto ri = _viToHi[r.vi];
            update(li, ri + 1);
            auto to = _parents[_leaders[r.vi]];
            r = {_depths[_leaders[to]], to};
        }
    }
};

struct State {
    ll cost{};
    State() {
    }
};

using S = struct {
    ll v;
    int len;
};
S op(S a, S b) {
    return {
        a.v + b.v,
        a.len + b.len,
    };
}
S e() {
    return {};
}
using F = ll;
S mapping(F f, S s) {
    return {
        s.v + f * s.len,
        s.len,
    };
}
F composition(F f, F g) {
    return f + g;
}
F id() {
    return 0;
}

using ST = atcoder::lazy_segtree<S, op, e, F, mapping, composition, id>;

int main() {
    ios::sync_with_stdio(false);
    std::cin.tie(nullptr);

    int N;
    cin >> N;
    HLD hld(N);
    rep(i, N - 1) {
        int u, v;
        cin >> u >> v;
        u--;
        v--;
        hld.AddEdge(u, v);
    }
    hld.Build();
    ST st(N);
    rep(i, N) {
        st.set(i, {1, 1});
    }

    int Q;
    cin >> Q;
    ll ans = 0;
    rep(i, Q) {
        int A, B;
        cin >> A >> B;
        A--;
        B--;
        P("A, B: %d, %d\n", A, B);
        auto q = hld.Query<State>(
            A, B,
            [&](int l, int r, const State &s) {
                auto v = st.prod(l, r);
                P("  Query: %d %d %lld\n", l, r, v);
                auto ans = s;
                ans.cost += v.v;
                return ans;
            },
            [](const State &a, const State &b) {
                auto ans = a;
                ans.cost += b.cost;
                return ans;
            });
        P("  q.cost: %lld\n", q.cost);
        ans += q.cost;
        hld.Update(A, B, [&](int l, int r) {
            P("  Apply: %d %d\n", l, r);
            st.apply(l, r, 1);
        });
        P("  st: ");
        rep(i, N) {
            P("%lld ", st.get(i));
        }
        P("\n");
    }
    cout << ans << endl;
    return 0;
}
0