13510번: 트리와 쿼리 1

https://www.acmicpc.net/problem/13510

HLD를 구현한다.

시간복잡도는 $O((n+m)\lg n)$

#include<cstdio>
#include<vector>
#include<algorithm>
using namespace std;
int last[100001], idx[100001], cnt[100001], par[100001], tree[400000], c;
int n, m, u[100000], v[100000], w[100000];
vector<int> adj[100001];
void count(int h, int p) {
    for (int it : adj[h]) if (it^p) count(it, h), cnt[h] += cnt[it];
    cnt[h]++;
}
void build(int h, int p) {
    int t = 0;
    for (int it : adj[h]) if (it^p && cnt[t] < cnt[it]) t = it;
    for (int it : adj[h]) if (it^p && it^t) build(it, h);
    if (!last[h]) last[h] = h;
    if (t) last[t] = last[h], build(t, h);
    par[h] = p;
    idx[h] = ++c;
}
void update(int h, int l, int r, int g, int x) {
    if (r < g || g < l) return;
    if (l == r) tree[h] = x;
    else {
        update(h * 2 + 1, l, (l + r) / 2, g, x);
        update(h * 2 + 2, (l + r) / 2 + 1, r, g, x);
        tree[h] = max(tree[h * 2 + 1], tree[h * 2 + 2]);
    }
}
int query(int h, int l, int r, int gl, int gr) {
    if (gr < l || r < gl) return 0;
    if (gl <= l&&r <= gr) return tree[h];
    return max(query(h * 2 + 1, l, (l + r) / 2, gl, gr), query(h * 2 + 2, (l + r) / 2 + 1, r, gl, gr));
}
int lca(int x, int y) {
    int ret = 0;
    while (last[x] ^ last[y]) {
        if (cnt[last[x]] > cnt[last[y]]) swap(x, y);
        ret = max(ret, query(0, 1, n, idx[x], idx[last[x]]));
        x = par[last[x]];
    }
    if (cnt[x] > cnt[y]) swap(x, y);
    return max(ret, query(0, 1, n, idx[x], idx[y] - 1));
}
int main() {
    scanf("%d", &n);
    for (int i = 1; i < n; i++) {
        scanf("%d%d%d", u + i, v + i, w + i);
        adj[u[i]].push_back(v[i]);
        adj[v[i]].push_back(u[i]);
    }
    count(1, 0);
    build(1, 0);
    for (int i = 1; i < n; i++) {
        if (par[v[i]] == u[i]) swap(u[i], v[i]);
        update(0, 1, n, idx[u[i]], w[i]);
    }
    scanf("%d", &m);
    for (int i = 0, q, x, y; i < m; i++) {
        scanf("%d%d%d", &q, &x, &y);
        if (q == 1) update(0, 1, n, idx[u[x]], y);
        else printf("%d\n", lca(x, y));
    }
    return 0;
}