dsu on tree 最后一个点(#28)TLE 求助
查看原帖
dsu on tree 最后一个点(#28)TLE 求助
95145
XUQING楼主2022/3/30 21:57

不知道最后一个点是什么类型的数据,照着题解对了下流程也没感觉哪里有问题

#include <iostream>
#include <vector>

#define FOR(i,a,b) for(int i = a; i <= b; i++)
typedef long long ll;
using namespace std;

const int N = 1e5+10;

ll c[N];
vector<int> G[2*N];

// 初始化重子树,dfn
int siz[N], hson[N];
void dfs1(int p, int fa) {
    siz[p] = 1;
    hson[p] = -1;
    for(auto son : G[p]) {
        if(son == fa) continue;
        dfs1(son, p);
        siz[p] += siz[son];
        if(hson[p] == -1 || siz[p] > siz[hson[p]]) {
            hson[p] = son;
        }
    }
}
int dfn[N], dfnR[N], idx[N];
int dfnCnt;
void dfs2(int p, int fa, int tp) {
    dfn[p] = ++dfnCnt;
    idx[dfnCnt] = p;
    if(hson[p] != -1) {
        dfs2(hson[p], p, tp);
        for(auto son : G[p]) {
            if (son != hson[p] && son != fa) {
                dfs2(son, p, son);
            }
        }
    }
    dfnR[p] = dfnCnt;
}

// 颜色个数,答案
ll cnt[N], ans[N];
// 主导颜色出现次数,目前答案
ll maxCnt, nowAns;
ll lastMaxCnt, lastNowAns;
void add(ll col) {
    cnt[col]++;
    if(cnt[col] > maxCnt) {
        maxCnt = cnt[col];
        nowAns = col;
    } else if(cnt[col] == maxCnt) {
        nowAns += col;
    }
}

void dsu(int p, int fa, bool keep) {
    lastMaxCnt = maxCnt; lastNowAns = nowAns;

    for(auto son : G[p]) {  // 递归计算轻子树答案 不保留
        if(son == fa || son == hson[p]) continue;
        dsu(son, p, false);
    }
    // 计算重子树答案 保留
    if(hson[p] != -1) dsu(hson[p], p, true);
    for(auto son : G[p]) {  // 再次遍历轻子树 统计答案
        if(son == fa || son == hson[p]) continue;
        FOR(i, dfn[son], dfnR[son]) {
            add(c[idx[i]]);
        }
    }
    add(c[p]);
    ans[p] = nowAns;

    if(!keep) { // 撤销这次记录的答案
        maxCnt = lastMaxCnt; nowAns = lastNowAns;
        FOR(i, dfn[p], dfnR[p]) {
            cnt[c[idx[i]]]--;
        }
    }
}

int main() {
    freopen("a.in", "r", stdin);
    ios::sync_with_stdio(false);
    cin.tie(0);

    int n;
    cin >> n;
    FOR(i,1,n) cin >> c[i];

    FOR(i,1,n-1) {
        int u,v;
        cin >> u >> v;
        G[u].push_back(v);
        G[v].push_back(u);
    }

    dfs1(1, 0);
    dfs2(1 ,0 , 1);
    dsu(1, 0, true);

    FOR(i,1,n) cout << ans[i] << " ";

    return 0;
}
2022/3/30 21:57
加载中...