#229146 (C++14) No.235 めぐるはめぐる (5)

提出ソース
結果

問題	No.235 めぐるはめぐる (5)
ユーザー	ふーらくたる
提出日時	2018-01-12 05:46:33
言語	C++14 (gcc 12.3.0 + boost 1.83.0)
結果	AC
実行時間	701 ms / 10,000 ms
コード長	11,313 bytes
コンパイル時間	1,617 ms
コンパイル使用メモリ	132,068 KB
実行使用メモリ	73,132 KB
最終ジャッジ日時	2024-06-06 07:29:20
合計ジャッジ時間	5,394 ms
ジャッジサーバーID （参考情報）	judge5 / judge3
このコードへのチャレンジ
（要ログイン）
テストケース

テストケース表示
入力	結果	実行時間実行使用メモリ
testcase_00	AC	701 ms 57,540 KB
testcase_01	AC	511 ms 73,132 KB
testcase_02	AC	673 ms 60,968 KB
権限があれば一括ダウンロードができます
ソースコード

raw source code
#include <iostream>
#include <cstdio>
#include <cstring>
#include <string>
#include <stack>
#include <queue>
#include <functional>
#include <algorithm>
#include <map>
#include <unordered_map>
#include <set>
#include <unordered_set>
#include <vector>
#include <array>
#include <tuple>
#include <utility>
#include <numeric>
#include <iomanip>
#include <cctype>
#include <cmath>
#include <assert.h>
#include <cstdlib>
#include <list>
using namespace std;

#define repeat(i, x) for (long long i = 0; (i) < (long long)(x); (i)++)
#define rrepeat(i, x) for (long long i = (long long)((x) - 1); (i) >= 0; (i)--)
#define traverse(it, ctn) for (auto it = (ctn).begin(); (it) != (ctn).end(); (it)++)
#define rtraverse(it, ctn) for (auto it = (ctn).rbegin(); (it) != (ctn).rend(); (it)++)
#define enumerate(i, a, b) for (long long i = (long long)(a); (i) < (long long)(b); (i)++)

template<typename T> void chmax(T& a1, T a2) { a1 = std::max(a1, a2); }
template<typename T> void chmin(T& a1, T a2) { a1 = std::min(a1, a2); }

template<typename T1, typename T2> ostream& operator<<(ostream& os, const pair<T1, T2>& p) { return os << "(" << p.first << ", " << p.second << ")"; }
template<typename T> ostream& operator<<(ostream& os, const vector<T>& v) { os << "["; for (int i = 0; i < v.size(); i++) os << (i == 0 ? "" : ", ") << v[i]; os << "]"; return os; }
template<typename T1, typename T2> ostream& operator << (ostream& os, map<T1, T2>& mp) { os << "{"; for (auto it = mp.begin(); it != mp.end(); it++) { os << "(" << it->first << ": " << it->second << ")"; it++; if(it != mp.end()) os << ", "; it--; } os << "}"; return os; }
template<typename T> ostream& operator << (ostream& os, set<T>& st) { os << "{"; for (auto it = st.begin(); it != st.end(); it++) { os << *it; ++it; if(it != st.end()) os << ", "; it--; } os << "}"; return os; }

using int64 = long long;
using Graph = std::vector<std::vector<int>>;

const int64 MOD = 1e9 + 7;

struct Element {
    using T = int64;
    static inline constexpr T identity() { return 0; }
    static inline T op(const T& a, const T& b) { return (a + b) % MOD; }
};

struct Operator {
    using T = int64;
    using M = T;
    static inline constexpr M identity() { return 0; }
    static inline M op(const M& a, const M& b) { return (a + b) % MOD; }
    // mpをaに適用する
    static inline T apply(const M& mp, const T a, const int64 w) { return ((w * T(mp)) % MOD + a) % MOD; }
};

template <class Object, class Operator>
class LazyPropagationSegmentTree {
    int N;
    int h;
    using T = typename Object::T;
    using M = typename Operator::M;
    std::vector<T> t;
    std::vector<M> lazy;
    std::vector<int64> co;

    inline int lowest_pow_of_2(int n) {
        int res = 1;
        while (res < n) res <<= 1;
        return res;
    }

    // log2_X <= nを満たす最大のXを計算
    inline int log2(int n) {
        int res = 0;
        while (n >> (res + 1)) res++;
        return res;
    }

    inline void prop_to(int i) {
        t[i] = Object::op(t[2 * i], t[2 * i + 1]);
    }

    inline void eval(int i, int w) {
        if (i < N and lazy[i] != Operator::identity()) {
            t[i] = Operator::apply(lazy[i], t[i], w);

            if (2 * i < N) {
                // 伝搬
                lazy[2 * i] = Operator::op(lazy[2 * i], lazy[i]);
                lazy[2 * i + 1] = Operator::op(lazy[2 * i + 1], lazy[i]);
            } else if (i < N) {
                t[2 * i] = Operator::apply(lazy[i], t[2 * i], co[2 * i]);
                t[2 * i + 1] = Operator::apply(lazy[i], t[2 * i + 1], co[2 * i + 1]);
            }
            lazy[i] = Operator::identity();
        }
    }

    public:

    LazyPropagationSegmentTree() {}

    LazyPropagationSegmentTree(int n)
        : N(lowest_pow_of_2(n)), h(log2(N) + 1),
          t(2 * N, Object::identity()), lazy(N, Operator::identity()) { }

    template <class InputIt>
    LazyPropagationSegmentTree(InputIt first, InputIt last)
        : N(lowest_pow_of_2(std::distance(first, last))),
          h(log2(N) + 1), 
          t(2 * N, Object::identity()),
          lazy(N, Operator::identity()) {
        std::copy(first, last, t.begin() + N);

        for (int i = N - 1; i > 0; i--) prop_to(i);
    }

    template <class InputIt>
    void build(InputIt first, InputIt last, std::vector<int64>& c) {
        N = lowest_pow_of_2(std::distance(first, last));
        h = log2(N) + 1;
        t = std::vector<T>(2 * N, Object::identity());
        lazy = std::vector<M>(N, Operator::identity());
        co = std::vector<int64>(2 * N, 0);

        std::copy(first, last, t.begin() + N);
        std::copy(c.begin(), c.end(), co.begin() + N);

        for (int i = N - 1; i > 0; i--) prop_to(i);
        for (int i = N - 1; i > 0; i--) co[i] = (co[2 * i] + co[2 * i + 1]) % MOD;
    }

    inline void update(int l, int r, M mp) {
        update(l, r, mp, 1, 0, N);
    }

    inline void update(int l, int r, M mp, int id, int nodel, int noder) {
        // 範囲外
        eval(id, co[id]);   // 演算の順序を守るためさきに伝播させる
        if (noder <= l or r <= nodel) return;

        if (id >= N) {  // 葉
            t[id] = Operator::apply(mp, t[id], co[id]);
            return;
        }

        if (l <= nodel and noder <= r) {
            lazy[id] = Operator::op(lazy[id], mp);
            eval(id, co[id]);
        } else {
            update(l, r, mp, 2 * id, nodel, (nodel + noder) / 2);
            update(l, r, mp, 2 * id + 1, (nodel + noder) / 2, noder);
            t[id] = Object::op(t[2 * id], t[2 * id + 1]);
        }
    }

    T get(int i) {
        i += N;
        for (int j = (h - 1); j > 0; j--) eval(i >> j, co[i >> j]);
        return t[i];
    }

    T query(int l, int r) {
        return query(l, r, 1, 0, N);
    }

    T query(int l, int r, int id, int nodel, int noder) {
        // 範囲外
        eval(id, co[id]);        
        if (noder <= l or r <= nodel) return Object::identity();

        // 完全に含まれる
        if (l <= nodel and noder <= r) return t[id];

        T resl = query(l, r, 2 * id, nodel, (nodel + noder) / 2);
        T resr = query(l, r, 2 * id + 1, (nodel + noder) / 2, noder);
        return Object::op(resl, resr);
    }
};

class HLDecomposition {
    using Graph = std::vector<std::vector<int>>;

    const int N;
    const int root;
    const Graph T;
    std::vector<std::vector<int>> chains;
    std::vector<int> depth;
    std::vector<int> subsize;
    std::vector<int> parent;
    std::vector<int> chain_id;
    std::vector<int> next;
    std::vector<int> at;
    std::vector<LazyPropagationSegmentTree<Element, Operator>> segtrees;

    // 各頂点の深さと部分木のサイズを計算する
    void setup() {
        stack<int> st;
        st.push(root);
        depth[root] = 0;

        while (not st.empty()) {
            int v = st.top(); st.pop();

            if (v >= 0) {
                st.push(~v);
                for (int to : T[v]) {
                    if (depth[to] >= 0) continue;
                    st.push(to);
                    parent[to] = v;
                    depth[to] = depth[v] + 1;
                }
            } else {
                v = ~v;
                subsize[v] = 1;
                for (int to : T[v]) {
                    if (parent[to] != v) continue;
                    subsize[v] += subsize[to];
                }
            }
        }
    }

    inline int head(int v) { return chains[chain_id[v]][0]; }
    inline int tail(int v) { return chains[chain_id[v]].back(); }
    inline int climb(int v) { return parent[head(v)]; }

    public:

    HLDecomposition(const Graph& t, int r)
        : N(std::distance(t.begin(), t.end())),
          root(r),
          T(t),
          chains(0),
          depth(N, -1),
          subsize(N, 0),
          parent(N, -1),
          chain_id(N, -1),
          next(N, -1),
          at(N, -1) {}

    void decompose(std::vector<int64>& S, std::vector<int64>& C) {
        setup();

        queue<int> Q;
        Q.push(root);
        while (not Q.empty()) {
            int v = Q.front(); Q.pop();

            if (chain_id[v] < 0) {
                chains.push_back(std::vector<int>());
                chain_id[v] = chains.size() - 1;
            } 
            int id = chain_id[v];
            at[v] = chains[id].size();
            chains[id].push_back(v);

            for (int to : T[v]) {
                if (parent[to] != v) continue;
                Q.push(to);
                if (next[v] < 0 or subsize[to] > subsize[next[v]]) {
                    next[v] = to;
                }
            }
            if (next[v] >= 0) chain_id[next[v]] = id;
        }

        // segtreesの構築
        segtrees.resize(chains.size());
        repeat (id, chains.size()) {
            std::vector<int64> elt(chains[id].size());
            std::vector<int64> co(chains[id].size());

            repeat (i, chains[id].size()) {
                int v = chains[id][i];
                elt[i] = S[v];
                co[i] = C[v];
            }

            segtrees[id].build(elt.begin(), elt.end(), co);
        }
    }

    int lca(int u, int v) {
        while (chain_id[u] != chain_id[v]) {
            if (depth[head(u)] > depth[head(v)]) u = climb(u);
            else v = climb(v);
        }
        return depth[u] > depth[v] ? v : u;
    }

    int get_parent(int v) { return parent[v]; }

    void update(int u, int v, int64 z) {
        while (chain_id[u] != chain_id[v]) {
            if (depth[head(u)] > depth[head(v)]) {
                segtrees[chain_id[u]].update(0, at[u] + 1, z);
                u = climb(u);
            } else {
                segtrees[chain_id[v]].update(0, at[v] + 1, z);
                v = climb(v);
            }
        }
        int hd, tl;
        if (depth[u] > depth[v]) { hd = at[v]; tl = at[u] + 1; }
        else { hd = at[u]; tl = at[v] + 1; }
        segtrees[chain_id[u]].update(hd, tl, z);
    }

    int64 query(int u, int v) {
        int64 res = 0;

        while (chain_id[u] != chain_id[v]) {
            if (depth[head(u)] > depth[head(v)]) {
                (res += segtrees[chain_id[u]].query(0, at[u] + 1)) %= MOD;
                u = climb(u);
            } else {
                (res += segtrees[chain_id[v]].query(0, at[v] + 1)) %= MOD;
                v = climb(v);
            }
        }
        int hd, tl;
        if (depth[u] > depth[v]) { hd = at[v]; tl = at[u] + 1; }
        else { hd = at[u]; tl = at[v] + 1; }
        return (res += segtrees[chain_id[u]].query(hd, tl)) %= MOD;
    }
};

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

    int N;
    cin >> N;

    vector<int64> S(N), C(N);

    repeat (i, N) cin >> S[i];
    repeat (i, N) cin >> C[i];

    Graph T(N);

    repeat (i, N - 1) {
        int A, B;
        cin >> A >> B;
        A--; B--;

        T[A].push_back(B);
        T[B].push_back(A);
    }

    HLDecomposition hl(T, 0);
    hl.decompose(S, C);

    int Q;
    cin >> Q;

    repeat (i, Q) {
        int com, X, Y;
        cin >> com >> X >> Y;
        X--; Y--;

        if (com == 0) {
            int64 Z;
            cin >> Z;
            hl.update(X, Y, Z);
        } else {
            cout << hl.query(X, Y) << '\n';
        }
    }

    return 0;
}
yukicoder

結果

テストケース

ソースコード
