#1040328 (C++23) No.235 めぐるはめぐる (5)

提出ソース
結果

問題	No.235 めぐるはめぐる (5)
ユーザー	Iroha_3856
提出日時	2025-02-01 20:51:59
言語	C++23 (gcc 13.3.0 + boost 1.87.0)
結果	AC
実行時間	2,640 ms / 10,000 ms
コード長	18,080 bytes
コンパイル時間	5,669 ms
コンパイル使用メモリ	306,824 KB
実行使用メモリ	53,624 KB
最終ジャッジ日時	2025-02-01 20:52:18
合計ジャッジ時間	16,651 ms
ジャッジサーバーID （参考情報）	judge2 / judge5
このコードへのチャレンジ
（要ログイン）
ファイルパターン	結果
other	AC * 3
権限があれば一括ダウンロードができます
ソースコード

raw source code
#include <bits/stdc++.h>
using namespace std;
#include <atcoder/lazysegtree>
#include <atcoder/modint>
using mint = atcoder::modint1000000007;

#define rep(i, l, r) for (int i = (int)(l); i<(int)(r); i++)

//遅延セグ木風のHLD(Heavy Light Decomposition, 重軽分解)
//パス更新を用いない場合、Fは適当に、mappingはそのままSを返す, compositionはidを返すなどにすればよい。
template<class S, auto op, auto e, 
         class F, auto mapping, auto composition, auto id>
struct HeavyLightDecomposition {
    static_assert(std::is_convertible_v<decltype(op), std::function<S(S, S)>>,
                  "op must work as S(S, S)");
    static_assert(std::is_convertible_v<decltype(e), std::function<S()>>,
                  "e must work as S()");
    static_assert(std::is_convertible_v<decltype(mapping), std::function<S(F, S)>>,
                  "mapping must work as S(F, S)");
    static_assert(std::is_convertible_v<decltype(composition), std::function<F(F, F)>>,
                  "composition must work as F(F, F)");
    static_assert(std::is_convertible_v<decltype(id), std::function<F()>>,
                  "id must work as F()");
    
    struct Edge {
        int to;
        S data;
        Edge(int t, S d) : to(t), data(d) {}
    };

    //Union-Find
    struct UnionFind {
        int _n;
        vector<int> par, siz;
        int connected_components;
        UnionFind() : UnionFind(0) {}
        //n頂点0辺のグラフを作る
        UnionFind(int n) : _n(n) {
            par.resize(n, -1);
            siz.resize(n, 1);
            connected_components = n;
        }
        //頂点xの連結成分の代表をreturnする
        int leader(int x) {
            if (par[x]==-1) return x;
            return par[x] = leader(par[x]);
        }
        //頂点xと頂点yが同じ連結成分に属しているか判定する
        bool same(int x, int y) {
            return leader(x)==leader(y);
        }
        //頂点xと頂点yを結ぶ辺を作る
        bool merge(int x, int y) {
            x = leader(x); y = leader(y);
            if (x==y) return false;
            //siz[x] > siz[y]を前提とする
            if (siz[x] < siz[y]) swap(x, y);
            par[y] = x;
            siz[x]+=siz[y];
            connected_components--;
            return true;
        }
        //頂点xの連結成分の頂点数をreturnする
        int size(int x) {
            return siz[leader(x)];
        }
        int connectedComponents() {
            return connected_components;
        }
        vector<vector<int>> groups() {
            vector<vector<int>> ret(_n);
            vector<int> group_size(_n, 0), leader_rem(_n);
            for (int i = 0; i < _n; i++) {
                leader_rem[i] = leader(i);
                group_size[leader_rem[i]]++;
            }
            for (int i = 0; i < _n; i++) ret[i].reserve(group_size[i]);
            for (int i = 0; i < _n; i++) {
                ret[leader_rem[i]].push_back(i);
            }
            ret.erase(remove_if(ret.begin(), ret.end(), [](const vector<int>& x) { return x.empty(); }), ret.end());
            return ret;
        }
    };

    int N, root;
    vector<int> roots;
    vector<vector<Edge>> G;
    UnionFind UF;

    //部分木サイズ, 行きがけ順, 帰りがけ順, その頂点から伸びるheavy edgeの行先, 親, そのheavy_edgeによるpathでの最上位頂点, inの逆配列
    vector<int> subtree_size, in, out, heavy_edge, par, head, where;
    vector<S> par_edge_data;
    vector<S> vertex_data;
    bool build_done;
    
    //遅延セグ木の変数宣言
    atcoder::lazy_segtree<S, op, e, F, mapping, composition, id> segv, sege;
    static S revop(S u, S d) { return op(d, u); }
    atcoder::lazy_segtree<S, revop, e, F, mapping, composition, id> revsegv, revsege;

    HeavyLightDecomposition(int n) : N(n), G(n), build_done(false), vertex_data(n, e()), par_edge_data(n), UF(n) {}
    HeavyLightDecomposition(vector<S>& A) : N((int)A.size()), G((int)A.size()), build_done(false), vertex_data(A), par_edge_data((int)A.size()), UF((int)A.size()) {}

    void addEdge(int u, int v, S x) {
        assert(0 <= u and u < N);
        assert(0 <= v and v < N);
        assert(!UF.same(u, v));
        G[u].push_back(Edge{v, x});
        G[v].push_back(Edge{u, x});
        UF.merge(u, v);
    }
    void addEdge(int u, int v) {
        assert(0 <= u and u < N);
        assert(0 <= v and v < N);
        assert(!UF.same(u, v));
        G[u].push_back(Edge{v, e()});
        G[v].push_back(Edge{u, e()});
        UF.merge(u, v);
    }
    void vectorResize() {
        subtree_size.resize(N, -1); in.resize(N); out.resize(N); heavy_edge.resize(N), par.resize(N), head.resize(N), where.resize(N);
    }
    bool isValid() {
        for (int i = 0; i < N; i++) if (subtree_size[i] == -1) return false;
        return true;
    }

    void buildTree(int _root = 0) {
        // cout << "ok : function build is executed." << endl;
        root = _root;
        vectorResize();
        // cout << "ok : resizing of vector id done." << endl;
        dfsSize(root, -1);
        // cout << "ok : function dfsSize is executed without RE." << endl;
        int t = 0;
        dfsHLD(root, -1, t);
        // cout << "ok : function dfsHLD is executed without RE." << endl;
        assert(isValid());
        // cout << "ok : function isValid is executed withoud RE." << endl;
        segSet();
        // cout << "ok : function segSet is executed without RE." << endl;
        build_done = true;
    }
    void buildForest() {
        vectorResize();
        int t = 0;
        vector<vector<int>> groups = UF.groups();
        for (vector<int> group : groups) {
            dfsSize(group[0], -1);
            dfsHLD(group[0], -1, t);
            roots.emplace_back(group[0]);
        }
        assert(isValid());
        segSet();
        build_done = true;
    }
    void buildForest(vector<int>& _roots) {
        roots = _roots;
        vectorResize();
        int t = 0;
        for (auto v : roots) {
            dfsSize(v, -1);
            dfsHLD(v, -1, t);
        }
        assert(isValid());
        segSet();
        build_done = true;
    } 
    
    //返り値は部分木のサイズ
    int dfsSize(int v, int p) {
        // cout << v << " : " << p << endl;
        par[v] = p;
        subtree_size[v] = 1;
        for (const Edge& ed : G[v]) {
            if (ed.to == p) continue;
            subtree_size[v] += dfsSize(ed.to, v);
        }
        // cout << v << " " << p << " : " << "done : subtree_size" << endl;
        int cmax = -1;
        heavy_edge[v] = -1;
        for (const Edge& ed : G[v]) {
            if (ed.to == p) continue;
            // cout << "v edge : " << ed.to << " " << "{" << ed.data.sum << ", " << ed.data.len << "}" << endl;
            par_edge_data[ed.to] = ed.data;
            if (subtree_size[ed.to] > cmax) {
                cmax = subtree_size[ed.to];
                heavy_edge[v] = ed.to;
            }
        }
        // cout << v << " " << p << " : done : successfully" << endl;
        return subtree_size[v];
    }
    void dfsHLD(int v, int p, int& t) {
        in[v] = t++;
        where[in[v]] = v;
        if (heavy_edge[v] != -1) {
            //heavy edgeが存在する（⇔葉ではない）
            head[heavy_edge[v]] = head[v];
            dfsHLD(heavy_edge[v], v, t);
        }
        for (const Edge& ed : G[v]) {
            if (ed.to == p or ed.to == heavy_edge[v]) continue;
            head[ed.to] = ed.to;
            dfsHLD(ed.to, v, t);
        }
        out[v] = t;
    }
    //セグメント木のセットアップ
    void segSet() {
        // cout << "======SUBTREE SIZE======\n";
        // rep(i, 0, N) cout << subtree_size[i] << " ";
        // cout << endl;
        // cout << "======HEAVY EDGE======\n";
        // rep(i, 0, N) cout << heavy_edge[i] << " ";
        // cout << endl;
        // cout << "======IN=====\n";
        // rep(i, 0, N) cout << in[i] << " ";
        // cout << endl;
        // cout << "======WHERE======\n";
        // rep(i, 0, N) cout << where[i] << " ";
        // cout << endl;
        // cout << "======HEAD======\n";
        // rep(i, 0, N) cout << head[i] << " ";
        // cout << endl;
        // cout << "======VERTEX DATA=====\n";
        // for (S x : vertex_data) cout << "{" << x.cost.val() << " " << x.sum.val() << "} ";
        // cout << endl;
        vector<S> segvdata(N);
        //上（小） -> 下（大）へと操作を行ったときどうなるか
        for (int i = 0; i < N; i++) {
            segvdata[i] = vertex_data[where[i]];
        }
        // cout << "======EDGE DATA=====\n";
        vector<S> segedata(N);
        //segedata[v] : par[v] <-> v へのデータが載る
        //上（小） -> 下（大）へと操作を行ったときどうなるか
        for (int i = 0; i < N; i++) {
            segedata[i] = (par[where[i]] != -1? par_edge_data[where[i]] : e());
            // cout << segedata[i].siz << " " << segedata[i].sum << ", ";
        }
        // cout << endl;
        
        segv = atcoder::lazy_segtree<S, op, e, F, mapping, composition, id>(segvdata);
        sege = atcoder::lazy_segtree<S, op, e, F, mapping, composition, id>(segedata);
        //下（大） -> 上（小）へと操作を行ったときどうなるか
        //載せるデータ自体は下 -> 上のものと変わらない。
        revsegv = atcoder::lazy_segtree<S, revop, e, F, mapping, composition, id>(segvdata);
        revsege = atcoder::lazy_segtree<S, revop, e, F, mapping, composition, id>(segedata);
    }

    //取得・更新クエリに答える関数

    //uからvへのパスについて、それに含まれる頂点の総積を計算する
    //O(log^2 N)
    S prodVertex(int u, int v) {
        assert(build_done);
        assert(UF.same(u, v));
        S retu = e(), retv = e();
        while(head[u] != head[v]) {
            if (in[u] > in[v]) {
                //uを上へ revsegを用いる
                retu = op(retu, revsegv.prod(in[head[u]], in[u]+1));
                u = par[head[u]];
            }
            else {
                //vを上へ segを用いる
                retv = op(segv.prod(in[head[v]], in[v]+1), retv);
                v = par[head[v]];
            }
        }
        assert(head[u] == head[v]);
        // cout << u << " " << v << " : " << in[u] << " " << in[v] << endl;
        if (in[u] < in[v]) {
            //u（上） -> v（下） segへクエリを投げる
            retv = op(segv.prod(in[u], in[v]+1), retv);
        }
        else {
            //v（上） -> u（下） revsegへクエリを投げる
            retu = op(retu, revsegv.prod(in[v], in[u]+1));
        }
        return op(retu, retv);
    }
    //uからvへのパスについて、それに含まれる辺の総積を計算する
    //O(log^2 N)
    S prodEdge(int u, int v) {
        // cout << "================" << endl;
        assert(build_done);
        assert(UF.same(u, v));
        S retu = e(), retv = e();
        while(head[u] != head[v]) {
            // cout << "(u, v) = " << u << ", " << v << endl;
            // cout << "(in[u], in[v]) = " << in[u] << ", " << in[v] << endl;
            if (in[u] > in[v]) {
                //uを上へ revsegを用いる
                // cout << "in[u] > in[v]" << endl;
                // cout << "head[u] = " << head[u] << endl;
                // cout << "in[head[u]], in[u] + 1 = " << in[head[u]] << " " << in[u]+1 << endl;
                retu = op(retu, revsege.prod(in[head[u]], in[u]+1));
                u = par[head[u]];
            }
            else {
                //vを上へ segを用いる
                // cout << "in[u] < in[v]" << endl;
                // cout << "head[v] = " << head[v] << endl;
                // cout << "in[head[v]], in[v] + 1 = " << in[head[v]] << " " << in[v]+1 << endl;
                retv = op(sege.prod(in[head[v]], in[v]+1), retv);
                v = par[head[v]];
            }
        }
        // cout << "end : " << endl;
        // cout << "(u, v) = " << u << ", " << v << endl;
        // cout << "(in[u], in[v]) = " << in[u] << ", " << in[v] << endl;
        // cout << "retu, retv = " << "(" << retu.siz << ", " << retu.sum << "), (" << retv.siz << ", " << retv.sum << ")" << endl;
        
        assert(head[u] == head[v]);
        // cout << in[u] << " " << in[v] << endl;
        if (in[u] < in[v]) {
            //u（上） -> v（下） segへクエリを投げる
            // cout << "A : " << sege.prod(in[u]+1, in[v]+1).sum << " " << sege.prod(in[u]+1, in[v]+1).len << endl;
            retv = op(sege.prod(in[u]+1, in[v]+1), retv);
        }
        else {
            //v（上） -> u（下） segへクエリを投げる
            // cout << "B : " << revsege.prod(in[v]+1, in[u]+1).sum << " " << revsege.prod(in[v]+1, in[u]+1).len << endl;
            retu = op(retu, revsege.prod(in[v]+1, in[u]+1));
        }
        // cout << retu.sum << " " << retu.len << "  " << retv.sum << " " << retv.len << endl;
        // cout << "======RETURN======\n";
        return op(retu, retv);
    }
    //頂点vのデータを返す
    int getVertex(int v) { assert(build_done); return segv.get(v); }

    //根がrootの木において、vとvの親頂点を結ぶ辺のデータを返す
    int getEdge(int v) { assert(build_done); return sege.get(v); }

    //uv間の辺に載っているデータを返す
    int getEdge(int u, int v) { 
        assert(build_done);
        assert(par[u] == v or par[v] == u);
        //par[v] = uとする
        if (par[u] == v) swap(u, v);
        return sege.get(v);
    }

    //uからvへのパスについて、そのパスに含まれる頂点にxの更新を行う
    //O(log^2 N)
    void applyVertex(int u, int v, F x) {
        assert(build_done);
        assert(UF.same(u, v));
        while(head[u] != head[v]) {
            if (in[u] > in[v]) {
                //uを上へ
                segv.apply(in[head[u]], in[u]+1, x);
                revsegv.apply(in[head[u]], in[u]+1, x);
                u = par[head[u]];
            }
            else {
                //vを上へ
                segv.apply(in[head[v]], in[v]+1, x);
                revsegv.apply(in[head[v]], in[v]+1, x);
                v = par[head[v]];
            }
        }
        assert(head[u] == head[v]);
        segv.apply(min(in[u], in[v]), max(in[u], in[v])+1, x);
        revsegv.apply(min(in[u], in[v]), max(in[u], in[v])+1, x);
    }

    //uからvへのパスについて、そのパスに含まれる辺にxの更新を行う
    //O(log^2 N)
    void applyEdge(int u, int v, F x) {
        assert(build_done);
        assert(UF.same(u, v));
        while(head[u] != head[v]) {
            if (in[u] > in[v]) {
                //uを上へ
                sege.apply(in[head[u]], in[u]+1, x);
                revsege.apply(in[head[u]], in[u]+1, x);
                u = par[head[u]];
            }
            else {
                //vを上へ
                sege.apply(in[head[v]], in[v]+1, x);
                revsege.apply(in[head[v]], in[v]+1, x);
                v = par[head[v]];
            }
        }
        assert(head[u] == head[v]);
        sege.apply(min(in[u], in[v])+1, max(in[u], in[v])+1, x);
        revsege.apply(min(in[u], in[v])+1, max(in[u], in[v])+1, x);
    }

    //頂点vのデータをxに更新する
    //O(log N)
    void setVertex(int v, S x) { segv.set(v, x); }

    //rootが根の木に置いて、vとvの親頂点を結ぶ辺のデータをxに更新する
    void setEdge(int v, S x) { sege.set(v, x); }

    //uv間の辺に載っているデータをxに更新する
    //もともとある辺にしか使えない
    //O(log N)
    void setEdge(int u, int v, S x) {
        //辺は元からあるはず
        assert(par[u] == v or par[v] == u);
        //par[v] = uとする。
        if (par[u] == v) swap(u, v);
        //par_edge_dataの更新は行わない（par_edge_dataはセグ木のための一時的なデータ）
        sege.set(v, x);
    }

    //u, vのLCAを返す
    //O(log N)
    int lca(int u, int v) {
        if (UF.same(u, v) == false) return -1;
        while(head[u] != head[v]) {
            if (in[u] < in[v]) v = par[head[v]];
            else u = par[head[u]];
        }
        return where[min(in[u], in[v])];
    }
    //vからrootへv遡った頂点を返す
    //存在しない場合は-1
    //O(log N)
    int la(int v, int k) {
        while(v != -1) {
            int u = head[v];
            if (in[u] <= in[v] - k) return where[in[v]-k];
            k -= (in[v]-in[u]+1);
            v = par[u]; 
        }
        return -1;
    }
};

using ll = long long;

//{コストの総和, S_xの総和}
struct S {
    mint cost, sum;
    S() : cost(0), sum(0) {}
    S(mint a, mint b) : cost(a), sum(b) {}
};
S op(S a, S b) {
    return S{a.cost+b.cost, a.sum+b.sum};
}
S e() { return S(); }
using F = mint;
S mapping(F f, S x) {
    return S{x.cost + x.sum * f, x.sum};
}
F composition(F f, F g) {
    return f + g;
}
F id() { return 0; }

int main() {
    // freopen("tree_input.txt", "r", stdin);
    // freopen("tree_output.txt", "w", stdout);
    int N; cin >> N;
    vector<int> s(N), c(N);
    rep(i, 0, N) cin >> s[i];
    rep(i, 0, N) cin >> c[i];
    vector<S> D(N);
    rep(i, 0, N) {
        D[i] = S{s[i], c[i]};
    }
    HeavyLightDecomposition<S, op, e, F, mapping, composition, id> HLD(D);
    rep(i, 0, N-1) {
        int u, v; cin >> u >> v;
        u--; v--;
        HLD.addEdge(u, v);
    }
    HLD.buildTree();
    int Q; cin >> Q;
    rep(i, 0, Q) {
        int q; cin >> q;
        if (q == 0) {
            int x, y, z; cin >> x >> y >> z;
            x--; y--;
            HLD.applyVertex(x, y, mint(z));
        }
        else {
            int u, v; cin >> u >> v;
            u--; v--;
            S ans = HLD.prodVertex(u, v);
            cout << ans.cost.val() << endl;
        }
    }
}
yukicoder

結果

ソースコード
