問題概要

辺にコストと色がついた $ N (\leq 10^5) $ 頂点の木について，以下のクエリに答えよ．クエリの回数は最大 $ 10^5 $ である．

色 $ x $ の辺のコストを全て $ y $ にしたときの頂点 $ u,v $ 間のコストを答えよ．

解法概要

LCA (Lowest Common Ancestor) を求めるパート
長さの変更を反映するパート

に分けて考える．

LCA はダブリングを使った方法などで求めると時間計算量 $ O(N \log N)$ で求められるので，$ u_i$と$ v_i$ のLCAが$ l_i$とわかっているものとして 2. について考える．

根から頂点$ i$までの経路について，その長さを$ d_i$，色$ x$ の辺の長さの和を$ d_{i,x}$，色$ x$の辺の本数を$ n_{i, x}$とすると，各クエリの答えは

$ (d_u+d_v-d_l) -(d_{u,x}+d_{v,x}-d_{l,x}) + y_i (n_{u_i,x_i}+n_{v_i,x_i}-2n_{l_i,x_i})$

$ d_{i,x}$ と$ n_{i,x}$ を全て求めようとすると時間計算量・空間計算量ともに$ O(N^2)$ となってしまい，間に合わない．これはクエリを先読みして必要な部分だけを map などに保管するようにすれば，時間計算量は$ O(N\log Q)$，空間計算量は$ O(Q)$に抑えることができる． (計算量の見積もりは怪しい．特に$ N$と$ Q$ を混同している可能性がある)

ソースコード

#include <iostream>
#include <vector>
#include <map>
#include <set>
#include <algorithm>

using namespace std;

using ll =  long long;

constexpr ll V_MAX = 100010;

vector<tuple<int, int, ll>> graph[V_MAX];
vector<ll> dists(V_MAX);
map<ll, int> colormp;
map<ll, ll> distmp;
set<int> required_colors[V_MAX];

class LowestCommonAncestor {
    public:
    int root, n, LOG_V_MAX = 20;
    vector<int> depth;
    vector<vector<int>> parent;

    LowestCommonAncestor(int root=0, int n=V_MAX): root(root), n(n) {
        depth.assign(n, -1);
        parent.assign(n, vector<int>(LOG_V_MAX, -1));
        dfs(root, -1, 0);
        for(int j=1;j<LOG_V_MAX;++j) {
            for(int i=0;i<n;++i) {
                if(parent[i][j-1] == -1) continue;
                parent[i][j] = parent[parent[i][j-1]][j-1];
            }
        }
    }

    void dfs(int node, int par, int d) {
        depth[node] = d;
        parent[node][0] = par;
        for(auto child: graph[node]) {
            int c = get<0>(child);
            if(c == par) continue;
            dfs(c, node, d+1);
        }
    }

    int get_lca(int u, int v) {
        if(depth[u] > depth[v]) swap(u, v);
        int depth_diff = depth[v] - depth[u];
        for(int j=0;j<LOG_V_MAX;++j) {
            if(depth_diff & (1 << j)) {
                v = parent[v][j];
            }
        }
        if(u == v) return u;
        for(int j=LOG_V_MAX-1;j>=0;--j) {
            if(parent[u][j] != parent[v][j]) {
                u = parent[u][j];
                v = parent[v][j];
            }
        }
        return parent[u][0];
    }
};

void dfs(int node, int par, ll dist, map<ll, int> &cmp, map<ll, ll> &dmp) {
    dists[node] = dist;
    for(int color: required_colors[node]) {
        colormp[V_MAX*node+color] = cmp[color];
        distmp[V_MAX*node+color] = dmp[color];
    }
    int chi, color; ll di;
    for(auto child: graph[node]) {
        tie(chi, color, di) = child;
        if(chi == par) continue;
        dist += di;
        ++cmp[color];
        dmp[color] += di;
        dfs(chi, node, dist, cmp, dmp);
        dist -= di;
        --cmp[color];
        dmp[color] -= di;
    }
}

int main() {
    int n,q;
    cin >> n >> q;
    int a,b,c; ll d;
    for(int i=0;i<n-1;++i){
        cin >> a >> b >> c >> d;
        --a; --b; --c;
        graph[a].push_back(make_tuple(b, c, d));
        graph[b].push_back(make_tuple(a, c, d));
    }

    LowestCommonAncestor lca = LowestCommonAncestor(0, n);

    vector<tuple<int,ll,int,int,int>> query(q);
    int u,v,l,x; ll y;
    for(int i=0;i<q;++i) {
        cin >> x >> y >> u >> v;
        --x; --u; --v;
        l = lca.get_lca(u, v);
        query[i] = make_tuple(x, y, u, v, l);
        required_colors[u].insert(x);
        required_colors[v].insert(x);
        required_colors[l].insert(x);
    }

    {
        map<ll, int> mp1;
        map<ll, ll> mp2;
        dfs(0, -1, 0, mp1, mp2);
    }

    for(int i=0;i<q;++i) {
        tie(x, y, u, v, l) = query[i];
        ll ans = (dists[u]+dists[v]-2*dists[l]) - (distmp[V_MAX*u+x]+distmp[V_MAX*v+x]-2*distmp[V_MAX*l+x]) + y*(colormp[V_MAX*u+x]+colormp[V_MAX*v+x]-2*colormp[V_MAX*l+x]);
        cout << ans << endl;
    }
}

所感

†プログラミングコンテスト† って感じで好き．

仮にまたKatsuT0shi でチーム戦に参加するとしてこういう問題が出たら私の担当だよな〜．ICPCは引退済みだけど重実装担当を自負する者としてこういう問題には強くありたい．

NoiminのNoise

AtCoder Beginner Contest 133 F - Colorful Tree

問題概要

解法概要

ソースコード

所感