2020牛客多校第一场B（虚树）

参考博客

#include<cstdio>
typedef long long ll;
const int N = 2e5 + 50;

int n, cnt, top, tot;
int c[N], mindiv[N], dep[N], w[N];
int  head[N], sta[N], lca_dep[N];
ll minn, ans[N];

struct Edge{
    int nex, to;
}e[N << 1];

inline int lowbit(int &x) { return x & -x; }
inline ll min(ll a, ll b) { return a < b ? a : b; }
inline void add(int a, int b) { e[++cnt] = {head[a], b};  head[a] = cnt; }

void add(int x){
    while(x <= n){
        ++c[x];
        x += lowbit(x);
    }
}

int ask(int x){
    int ans = 0;
    while(x){
        ans += c[x];
        x -= lowbit(x);
    }
    return ans;
}

void init(int n){
    cnt = 0;
    for(int i = 1; i <= n; ++i)  c[i] = w[i] = dep[i] = lca_dep[i] = ans[i] = 0;
}

void pre(int n){
    for(int i = 2; i <= n; ++i)
        for(int t = i; t <= n; t += i)
            if(!mindiv[t])  mindiv[t] = i;
}

void build(){
    tot = n;
    sta[top = 1] = 1, head[1] = 0;
    for(int i = 2, t; i <= n; ++i){
        t = i;  dep[i] = dep[i - 1] + 1;
        while(t != mindiv[t])  ++dep[i], t /= mindiv[t];
        lca_dep[i] = ask(n) - ask(t - 1);
        for(int t = i; t != 1; t /= mindiv[t])  add(mindiv[t]);
    }
    for(int i = 2; i <= n; ++i){
    	//printf("%d %d\n", lca_dep[i], sta[1]);
        while(top > 1 && dep[sta[top - 1]] >= lca_dep[i]){
            add(sta[top - 1], sta[top]);
            --top;
        }
        if(dep[sta[top]] != lca_dep[i]){
            dep[++tot] = lca_dep[i];
            head[tot] = 0;
            add(tot, sta[top]);
            sta[top] = tot;
        }
        head[i] = 0,  sta[++top] = i;
    }
    //printf("%d %d ", top, sta[1]);
    for(int i = 1; i < top; ++i)  add(sta[i], sta[i + 1]);
}

void dfs(int u){
	//printf("%d ", u);
    ans[1] += 1LL * w[u] * dep[u];
    for(int i = head[u]; i; i = e[i].nex){
        int to = e[i].to;
        dfs(to);
        w[u] += w[to];
    }
}

void dfs2(int u){
    for(int i = head[u]; i; i = e[i].nex){
        int to = e[i].to;
        ans[to] = ans[u] + (w[1] - 2LL * w[to]) * (dep[to] - dep[u]);
        dfs2(to);
    }
}

int main(){
    pre(1e5);
    while(scanf("%d", &n) != EOF){
        for(int i = 1; i <= n; ++i)  scanf("%d", &w[i]);
        minn = 1e15;
        build();
        dfs(1);  dfs2(1);
        for(int i = 1; i <= tot; ++i)  minn = min(ans[i], minn);
        printf("%lld\n", minn);
        init(tot);
    }
    return 0;
}