P3714 [BJOI2017] 树的难题

题意

给一颗 \(n\) 个点的树，边有颜色，颜色共有 \(m\) 种，每个颜色都有权值，定义一条路径的长度为路径上的边的颜色组成的序列中同每个颜色段的颜色的权值和。求长度在 \([l,r]\) 的最大路径的权值和。
\(n,m\le 2\times10^5\)。

思路

与树路径统计有关，考虑点分治。
对于一个分治中心，把祂的出边按照颜色排序，这样同色的出边就排到一起了。按顺序枚举出边，开两个线段树，一个维护从分治中心向下走的第一条边与当前的边同色的路径信息，另一个维护颜色不一样的边，线段树的下标表示路径长度。一条路径的权值即为从分治中心出发两条路径的权值之和，如果祂们从分治中心向下第一条边颜色相同，哪一种颜色的权值会被算两次，需要减去。
时间复杂度 \(\mathcal O(nlog^2n)\)。

代码

/*
Luogu P3714 [BJOI2017] 树的难题
2026-04-16
*/
#include<bits/stdc++.h>
using namespace std;
namespace IO{
    template<typename T>
    inline void read(T&x){
        x=0;char c=getchar();bool f=0;
        while(!isdigit(c)) c=='-'?f=1:0,c=getchar();
        while(isdigit(c)) x=x*10+c-'0',c=getchar();
        f?x=-x:0;
    }
    template<typename T>
    inline void write(T x){
        if(x==0){putchar('0');return ;}
        x<0?x=-x,putchar('-'):0;short st[50],top=0;
        while(x) st[++top]=x%10,x/=10;
        while(top) putchar(st[top--]+'0');
    }
    inline void read(char&c){c=getchar();while(isspace(c)) c=getchar();}
    inline void write(char c){putchar(c);}
    inline void read(string&s){s.clear();char c;read(c);while(!isspace(c)&&~c) s+=c,c=getchar();}
    inline void write(string s){for(int i=0,len=s.size();i<len;i++) putchar(s[i]);}
    template<typename T>inline void write(T*x){while(*x) putchar(*(x++));}
    template<typename T,typename...T2> inline void read(T&x,T2&...y){read(x),read(y...);}
    template<typename T,typename...T2> inline void write(const T x,const T2...y){write(x),putchar(' '),write(y...),sizeof...(y)==1?putchar('\n'):0;}
}using namespace IO;
const int maxn=200010,inf=2000000000;
vector<pair<int,int>>e[maxn];
int n,m,l,r,c[maxn],son[maxn],d,sz[maxn],rt,ans=-inf,val;
bool flag[maxn];
class Segment_Tree{
private:
    int t[maxn*4];
    void update(int u,int l,int r,int d,int z){
        if(l>d||r<d) return ;
        if(l==r){t[u]=max(z,t[u]);return ;}
        int mid=l+r>>1;
        update(u<<1,l,mid,d,z),update(u<<1|1,mid+1,r,d,z);
        t[u]=max(t[u<<1],t[u<<1|1]);
    }
    int query(int u,int l,int r,int ll,int rr){
        if(l>rr||r<ll) return -inf;
        if(ll<=l&&r<=rr) return t[u];
        int mid=l+r>>1;
        return max(query(u<<1,l,mid,ll,rr),query(u<<1|1,mid+1,r,ll,rr));
    }
    void clear(int u,int l,int r,int d){
        if(l>d||r<d) return ;
        if(l==r){t[u]=-inf;return ;}
        int mid=l+r>>1;
        clear(u<<1,l,mid,d),clear(u<<1|1,mid+1,r,d);
        t[u]=max(t[u<<1],t[u<<1|1]);
    }
public:
    Segment_Tree(){for(int i=0;i<maxn*4;i++) t[i]=-inf;}
    void update(int d,int z){update(1,0,n,d,z);}
    void clear(int d){clear(1,0,n,d);}
    int query(int l,int r){return query(1,0,n,l,r);}
}t1,t2;
void get_rt(int u,int fa=0){
    sz[u]=1,son[u]=0;
    for(auto[v,w]:e[u]){
        if(v==fa||flag[v]) continue;
        get_rt(v,u);
        sz[u]+=sz[v],son[u]=max(son[u],sz[v]);
    }
    son[u]=max(son[u],d-sz[u]);
    if(son[u]<son[rt]) rt=u;
}
void update2(int u,int col,int z,int deep,int fa=0){
    if(deep>r) return ;
    t2.update(deep,z);
    for(auto[v,w]:e[u]){
        if(v==fa||flag[v]) continue;
        if(w==col) update2(v,w,z,deep+1,u);
        else update2(v,w,z+c[w],deep+1,u);
    }
}
void clear1(int u,int deep,int fa=0){
    if(deep>r) return ;
    t1.clear(deep);
    for(auto[v,w]:e[u]){
        if(v==fa||flag[v]) continue;
        clear1(v,deep+1,u);
    }
}
void clear2(int u,int deep,int fa=0){
    if(deep>r) return ;
    t2.clear(deep);
    for(auto[v,w]:e[u]){
        if(v==fa||flag[v]) continue;
        clear2(v,deep+1,u);
    }
}
void update1(int u,int col,int z,int deep,int fa=0){
    if(deep>r) return ;
    t1.update(deep,z);
    for(auto[v,w]:e[u]){
        if(v==fa||flag[v]) continue;
        if(w==col) update1(v,w,z,deep+1,u);
        else update1(v,w,z+c[w],deep+1,u);
    }
}
void query(int u,int col,int z,int deep,int fa=0){
    if(deep>r) return ;
    ans=max(ans,t2.query(l-deep,r-deep)+z-c[val]);
    ans=max(ans,t1.query(l-deep,r-deep)+z);
    for(auto[v,w]:e[u]){
        if(v==fa||flag[v]) continue;
        if(w==col) query(v,w,z,deep+1,u);
        else query(v,w,z+c[w],deep+1,u);
    }
}
void dfs(int u,int fa=0){
    flag[u]=1;
    int lt=0;
    vector<pair<int,int>>rl;
    for(auto x:e[u]) if(!flag[x.first]) rl.push_back(x);
    if(rl.size()){
        t1.update(0,0),val=rl[0].second;
        for(int i=0;i<rl.size();i++){
            auto[v,w]=rl[i];
            if(i&&w!=rl[i-1].second){
                for(int j=lt;j<i;j++){
                    auto[vv,ww]=rl[j];
                    clear2(vv,1);
                    update1(vv,ww,c[ww],1);
                }
                lt=i,val=w;
            }
            query(v,w,c[w],1);
            update2(v,w,c[w],1);
        }
        clear1(u,0),clear2(u,0);
    }
    for(auto[v,w]:e[u]){
        if(flag[v]) continue;
        rt=0,d=sz[v];
        get_rt(v);
        dfs(rt);
    }
}
signed main(){
    read(n,m,l,r);
    for(int i=1;i<=m;i++) read(c[i]);
    for(int i=1;i<n;i++){
        int u,v,c;read(u,v,c);
        e[u].push_back({v,c}),e[v].push_back({u,c});
    }
    for(int i=1;i<=n;i++) sort(e[i].begin(),e[i].end(),[](pair<int,int>a,pair<int,int>b){return a.second<b.second;});
    son[0]=d=n;
    get_rt(1);
    dfs(rt);
    write(ans);
    return 0;
}