[ICPC 2018 Asia Yokohama Onsite J]Colorful Tree

Link

Solution

年轻的时候见这种题比较多。。。
需要维护每种颜色的点的集合,和每种颜色点的LCA。
点的集合按照DFS序来存储,怎么维护想想就明白了。

Code

const int XN=1e5+11,XL=20;

#define log2 _log2

int log2[XN];
std::vector<int> G[XN];
int dfn[XN],dfs[XN],dfc,fa[XN],dep[XN];
int par[XN][XL];

void DFS(int pos) {
    par[pos][0]=fa[pos];
    for(int b=1;b<=log2[dep[pos]];++b)
        par[pos][b]=par[par[pos][b-1]][b-1];

    dfs[dfn[pos]=++dfc]=pos;

    for(auto u : G[pos]) if(u!=fa[pos]) {
        fa[u]=pos;
        dep[u]=dep[pos]+1;
        DFS(u);
    }
}

int LCA(int u,int v) {
    if(u==-1 || v==-1)
        return u==-1?v:u;
    else {
        if(dep[u]<dep[v])
            std::swap(u,v);
        for(int len=dep[u]-dep[v],b=log2[len];b>=0;--b)
            if(len>>b&1)
                u=par[u][b];
        if(u!=v) {
            for(int b=log2[dep[u]];b>=0;--b)
                if(par[u][b]!=par[v][b]) {
                    u=par[u][b];
                    v=par[v][b];
                }
            u=par[u][0];
        }
        assert(u);
        return u;
    }
}

int Dist(int u,int v,int lca=0) {
    return dep[u]+dep[v]-2*dep[lca?lca:LCA(u,v)];
}

int Ans[XN],subLCA[XN];
std::set<int> S[XN];

int Up(int pos,int c) {
    int up=0;
    auto iter=S[c].lower_bound(dfn[pos]);
    if(iter!=S[c].begin()) {
        auto pred=(--iter)++;
        int lca=LCA(dfs[*pred],pos);
        if(dep[up]<dep[lca])
            up=lca;
    }
    auto succ=++iter;
    if(succ!=S[c].end()) {
        int lca=LCA(dfs[*succ],pos);
        if(dep[up]<dep[lca])
            up=lca;
    }
    return up;
}

void Add(int pos,int c) {
    S[c].insert(dfn[pos]);
    if(subLCA[c]==-1)
        subLCA[c]=pos;
    else {
        int up=Up(pos,c);
        if(dep[up]<dep[subLCA[c]]) {
            Ans[c]+=Dist(subLCA[c],pos,up);
            subLCA[c]=up;
        } else
            Ans[c]+=dep[pos]-dep[up];
    }
}

void Delete(int pos,int c) {
    if(S[c].size()==1) {
        subLCA[c]=-1;
        S[c].erase(dfn[pos]);
    } else {
        int up=Up(pos,c);       
        S[c].erase(dfn[pos]);
        if(up==subLCA[c]) {
            int lca=S[c].size()==1?dfs[*S[c].begin()]:LCA(dfs[*S[c].begin()],dfs[*S[c].rbegin()]);
            Ans[c]-=Dist(lca,pos);
            subLCA[c]=lca;
        } else {
            Ans[c]-=dep[pos]-dep[up];
        }
    }
}

int main() {
    freopen("input","r",stdin);
    std::ios::sync_with_stdio(0);
    std::cin.tie(0);
    int n;std::cin>>n;
    log2[0]=-1;
    for(int i=1;i<=n-1;++i) {
        int x,y;std::cin>>x>>y;
        G[x].push_back(y);
        G[y].push_back(x);
        //std::cerr<<i<<':';
        log2[i]=log2[i>>1]+1;
    }
    dep[1]=1;
    DFS(1);

    memset(subLCA,-1,sizeof(subLCA));

    static int c[XN];
    for(int i=1;i<=n;++i) {
        std::cin>>c[i];
        Add(i,c[i]);
    }
    int m;std::cin>>m;
    for(int q=1;q<=m;++q) {
        char op;int x,y;
        std::cin>>op>>x;
        if(op=='Q') {
            if(S[x].empty())
                std::cout<<-1<<'\n';
            else
                std::cout<<Ans[x]<<'\n';
        } else {
            std::cin>>y;
            Delete(x,c[x]);
            Add(x,c[x]=y);
        }
    }
    return 0;
}

发表评论

电子邮件地址不会被公开。 必填项已用*标注