SDOI2015 寻宝游戏

#include <cstdio>

#include <set>

#define LOG(FMT...) // fprintf(stderr, FMT)

using namespace std;

typedef long long ll;
typedef set<int>::iterator iter;

struct edge {
    int v, w;
    edge* next;
};

const int N = 100010;

int n, m;
edge* g[N];
int vis[N], dep[N], rev[N];
int prt[N][20];
ll sum[N][20];

set<int> st;

void add_edge(int u, int v, int w);
void dfs(int u);
ll get_dist(int u, int v);
iter my_next(iter it);
iter my_prev(iter it);

int main() {
    int u, v, w;
    iter it;
    ll ans = 0;
    scanf("%d%d", &n, &m);
    for (int rep = 1; rep < n; ++rep) {
        scanf("%d%d%d", &u, &v, &w);
        add_edge(u, v, w);
        add_edge(v, u, w);
    }
    prt[1][0] = 1;
    dfs(1);
    while (m--) {
        scanf("%d", &u);
        if ((it = st.find(vis[u])) == st.end()) {
            it = st.insert(vis[u]).first;
            ans += get_dist(u, rev[*my_next(it)])
                 + get_dist(u, rev[*my_prev(it)])
                 - get_dist(rev[*my_next(it)], rev[*my_prev(it)]);
        } else {
            ans += get_dist(rev[*my_prev(it)], rev[*my_next(it)])
                 - get_dist(u, rev[*my_next(it)])
                 - get_dist(u, rev[*my_prev(it)]);
            LOG("#%d\n", *st.end());
            st.erase(it);
        }
        printf("%lld\n", ans);
    }
    return 0;
}

inline iter my_next(iter it) {
    return ++it == st.end() ? st.begin() : it;
}

inline iter my_prev(iter it) {
    --it;
    if (it == st.end() && st.size() > 1)
        --it;
    return it;
}

inline void add_edge(int u, int v, int w) {
    static edge pool[N << 1];
    static edge* p = pool;
    p->v = v;
    p->w = w;
    p->next = g[u];
    g[u] = p;
    ++p;
}

void dfs(int u) {
    static int time;
    vis[u] = ++time;
    rev[time] = u;
    for (int i = 1; i < 20; ++i) {
        prt[u][i] = prt[prt[u][i - 1]][i - 1];
        sum[u][i] = sum[u][i - 1] + sum[prt[u][i - 1]][i - 1];
    }
    for (edge* p = g[u]; p; p = p->next)
        if (!vis[p->v]) {
            prt[p->v][0] = u;
            sum[p->v][0] = p->w;
            dep[p->v] = dep[u] + 1;
            dfs(p->v);
        }
}

inline ll get_dist(int u, int v) {
    ll ret = 0;
    LOG("%d %d: ", u, v);
    if (dep[u] < dep[v])
        swap(u, v);
    for (int i = 19; ~i; --i)
        if (dep[prt[u][i]] >= dep[v]) {
            ret += sum[u][i];
            u = prt[u][i];
        }
    if (u == v) {
        LOG("%lld\n", ret);
        return ret;
    }
    for (int i = 19; ~i; --i)
        if (prt[u][i] != prt[v][i]) {
            ret += sum[u][i] + sum[v][i];
            u = prt[u][i];
            v = prt[v][i];
        }
    ret += sum[u][0] + sum[v][0];
    LOG("%lld\n", ret);
    return ret;
}