結果

問題 No.386 貪欲な領主
ユーザー kimiyukikimiyuki
提出日時 2016-07-02 00:08:13
言語 C++14
(gcc 13.3.0 + boost 1.87.0)
結果
AC  
実行時間 999 ms / 2,000 ms
コード長 7,012 bytes
コンパイル時間 4,241 ms
コンパイル使用メモリ 104,512 KB
実行使用メモリ 38,008 KB
最終ジャッジ日時 2024-10-12 19:17:01
合計ジャッジ時間 7,676 ms
ジャッジサーバーID
（参考情報） 		judge5 / judge2
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
other AC * 16
権限があれば一括ダウンロードができます

ソースコード

diff # 

#include <cstdio>
#include <vector>
#include <algorithm>
#include <array>
#include <set>
#include <map>
#include <queue>
#include <tuple>
#include <unordered_set>
#include <unordered_map>
#include <functional>
#include <cmath>
#include <cassert>
#define repeat(i,n) for (int i = 0; (i) < (n); ++(i))
#define repeat_from(i,m,n) for (int i = (m); (i) < (n); ++(i))
#define repeat_reverse(i,n) for (int i = (n)-1; (i) >= 0; --(i))
#define repeat_from_reverse(i,m,n) for (int i = (n)-1; (i) >= (m); --(i))
#define whole(f,x,...) ([&](decltype((x)) y) { return (f)(begin(y), end(y), ## __VA_ARGS__); })(x)
typedef long long ll;
using namespace std;
template <class T> void setmax(T & a, T const & b) { if (a < b) a = b; }
template <class T> void setmin(T & a, T const & b) { if (b < a) a = b; }
template <typename T, typename X> auto vectors(T a, X x) { return vector<T>(x, a); }
template <typename T, typename X, typename Y, typename... Zs> auto vectors(T a, X x, Y y, Zs... zs) { auto cont = vectors(a, y, zs...); return vector<decltype(cont)>(x, cont); }

struct heavy_light_decomposition_t {
    int n; // |V'|
    vector<int> a; // V ->> V' epic
    vector<vector<int> > path; // V' -> V*, bottom to top order, disjoint union of codomain matchs V
    vector<map<int,int> > pfind; // V' * V -> int, find in path
    vector<int> parent; // V' -> V
    heavy_light_decomposition_t(int v, vector<vector<int> > const & g) {
        n = 0;
        a.resize(g.size());
        dfs(v, -1, g);
    }
    int dfs(int v, int p, vector<vector<int> > const & g) {
        int heavy_node = -1;
        int heavy_size = 0;
        int desc_size = 1;
        for (int w : g[v]) if (w != p) {
            int size = dfs(w, v, g);
            desc_size += size;
            if (heavy_size < size) {
                heavy_node = w;
                heavy_size = size;
            }
        }
        if (heavy_node == -1) {
            a[v] = n;
            n += 1;
            path.emplace_back();
            path.back().push_back(v);
            pfind.emplace_back();
            pfind.back()[v] = 0;
            parent.push_back(p);
        } else {
            int i = a[heavy_node];
            a[v] = i;
            pfind[i][v] = path[i].size();
            path[i].push_back(v);
            parent[i] = p;
        }
        return desc_size;
    }
};

struct lowest_common_ancestor_t {
    vector<vector<int> > a;
    vector<int> depth;
    lowest_common_ancestor_t(int v, vector<vector<int> > const & g) {
        int n = g.size();
        int l = 1 + floor(log2(n));
        a.resize(l);
        repeat (k,l) a[k].resize(n, -1);
        depth.resize(n);
        dfs(v, -1, 0, g, a[0], depth);
        repeat (k,l-1) {
            repeat (i,n) {
                if (a[k][i] != -1) {
                    a[k+1][i] = a[k][a[k][i]];
                }
            }
        }
    }
    static void dfs(int v, int p, int current_depth, vector<vector<int> > const & g, vector<int> & parent, vector<int> & depth) {
        parent[v] = p;
        depth[v] = current_depth;
        for (int w : g[v]) if (w != p) {
            dfs(w, v, current_depth + 1, g, parent, depth);
        }
    }
    int operator () (int x, int y) const {
        int l = a.size();
        if (depth[x] < depth[y]) swap(x,y);
        repeat_reverse (k,l) {
            if (a[k][x] != -1 and depth[a[k][x]] >= depth[y]) {
                x = a[k][x];
            }
        }
        assert (depth[x] == depth[y]);
        if (x == y) return x;
        repeat_reverse (k,l) {
            if (a[k][x] != a[k][y]) {
                x = a[k][x];
                y = a[k][y];
            }
        }
        assert (x != y);
        assert (a[0][x] == a[0][y]);
        return a[0][x];
    }
    int descendant_to (int x, int y) const {
        assert (depth[x] < depth[y]);
        int l = a.size();
        repeat_reverse (k,l) {
            if (a[k][y] != -1 and depth[a[k][y]] >= depth[x]+1) {
                y = a[k][y];
            }
        }
        assert (a[0][y] == x);
        return y;
    }
};

template <typename T>
struct segment_tree { // on monoid
    int n;
    vector<T> a;
    function<T (T,T)> append; // associative
    T unit;
    template <typename F>
    segment_tree(int a_n, T a_unit, F a_append) {
        n = pow(2,ceil(log2(a_n)));
        a.resize(2*n-1, a_unit);
        unit = a_unit;
        append = a_append;
    }
    void point_update(int i, T z) {
        a[i+n-1] = z;
        for (i = (i+n)/2; i > 0; i /= 2) {
            a[i-1] = append(a[2*i-1], a[2*i]);
        }
    }
    T range_concat(int l, int r) {
        return range_concat(0, 0, n, l, r);
    }
    T range_concat(int i, int il, int ir, int l, int r) {
        if (l <= il and ir <= r) {
            return a[i];
        } else if (ir <= l or r <= il) {
            return unit;
        } else {
            return append(
                    range_concat(2*i+1, il, (il+ir)/2, l, r),
                    range_concat(2*i+2, (il+ir)/2, ir, l, r));
        }
    }
};

template <typename T>
T path_concat(heavy_light_decomposition_t & hl, vector<segment_tree<T> > & sts, int v, int w) {
    auto append = sts.front().append;
    auto unit   = sts.front().unit;
    T acc = unit;
    int i = hl.a[v];
    if (hl.a[w] == i) {
        assert (hl.pfind[i][v] <= hl.pfind[i][w]); // v must be a descendant of w
        acc = append(acc, sts[i].range_concat(hl.pfind[i][v], hl.pfind[i][w]+1));
    } else {
        acc = append(acc, sts[i].range_concat(hl.pfind[i][v], hl.path[i].size()));
        acc = append(acc, path_concat(hl, sts, hl.parent[i], w));
    }
    return acc;
}

template <typename T>
T path_concat(heavy_light_decomposition_t & hl, lowest_common_ancestor_t & lca, vector<segment_tree<T> > & sts, int x, int y) {
    auto append = sts.front().append;
    auto unit   = sts.front().unit;
    int z = lca(x, y);
    T acc = unit;
    if (x != z) acc = append(acc, path_concat(hl, sts, x, lca.descendant_to(z, x)));
    if (y != z) acc = append(acc, path_concat(hl, sts, y, lca.descendant_to(z, y)));
    acc = append(acc, path_concat(hl, sts, z, z));
    return acc;
}

int main() {
    // input
    int n; scanf("%d", &n);
    vector<vector<int> > g(n);
    repeat (i,n-1) {
        int a, b; scanf("%d%d", &a, &b);
        g[a].push_back(b);
        g[b].push_back(a);
    }
    vector<int> u(n); repeat (i,n) scanf("%d", &u[i]);
    // prepare
    const int root = 0;
    heavy_light_decomposition_t hl(root, g);
    vector<segment_tree<ll> > sts;
    repeat (i,hl.n) {
        sts.emplace_back(hl.path[i].size(), 0, [](ll a, ll b) { return a + b; });
    }
    repeat (i,n) {
        int l = hl.a[i];
        sts[l].point_update(hl.pfind[l][i], u[i]);
    }
    lowest_common_ancestor_t lca(root, g);
    // run
    ll ans = 0;
    int m; scanf("%d", &m);
    repeat (i,m) {
        int a, b, c; scanf("%d%d%d", &a, &b, &c);
        ans += path_concat(hl, lca, sts, a, b) * c;
    }
    // output
    printf("%lld\n", ans);
    return 0;
}
0