题目描述

给你一棵 n 个点的无根树。

树上的每条边具有颜色。一共有 m 种颜色，编号为 1 到 m，第 i 种颜色的权值为 c[i]

对于一条树上的简单路径，路径上经过的所有边按顺序组成一个颜色序列，序列可以划分成若干个相同颜色段。
定义路径权值为颜色序列上每个同颜色段的颜色权值之和。

请你计算，经过边数在 l 到 r 之间的所有简单路径中，路径权值的最大值。

#include <bits/stdc++.h>

#define INF int(1e9)
#define getD(p) (p ? p->d : -INF)

using namespace std;

typedef long long LL;

const int N = 2e5 + 7;

int n, m, L, R;
// val表示所有颜色的权值 dep每个点的深度 dis存路径 col表示每个点的颜色
int val[N], dep[N], dis[N], col[N];
struct Edge{
    int v, c;
    Edge(int v = 0, int c = 0)
        :v(v), c(c){}
    bool operator<(const Edge& e)const{
        return c < e.c;
    }
};
vector<Edge> adj[N];
int ans = -2e9;

int siz[N], maxp[N], rt, sum;
bool vis[N];
void getrt(int u,int f){

    siz[u] = 1, maxp[u] = 0;
    for(int i = 0;i < adj[u].size();i ++ ){

        int v = adj[u][i].v;
        if(v == f || vis[v]) continue;
        getrt(v, u);
        siz[u] += siz[v];
        if(siz[v] > maxp[u]) maxp[u] = siz[v];
    }
    maxp[u] = max(maxp[u], sum - siz[u]);
    if(maxp[u] < maxp[rt]) rt = u;
}

// 用深度建立线段树(动态开点) 每一层维护一个最大值
struct Node{
    int d; // 维护最大值
    Node* ls, *rs;
    void update(){
        d = max(getD(ls), getD(rs));
    }
}tr[N * 50], *rt0, *rt1;
int top = 0;

// 当前节点 左区间 右区间 编号 要插入的值
void insert(Node* &now, int l, int r, int id, int x){
    if(!now){
        now = tr + ( ++ top);
        now->d = -INF;
        now->ls = now->rs = NULL;
    }

    if(l == r){
        now->d = max(now->d, x);
        return ;
    }

    int mid = l + r >> 1;

    if(id <= mid) insert(now->ls, l, mid, id, x);
    else insert(now->rs, mid + 1, r, id, x);

    now->update();
}
// 合并 p 和 q 两棵线段树
Node* merge(Node* p, Node* q, int l, int r){
    if(!p || !q) return p ? p : q;

    if(l == r){
        p->d = max(p->d, q->d);
        return p;
    }

    int mid = l + r >> 1;
    p->ls = merge(p->ls, q->ls, l, mid);
    p->rs = merge(p->rs, q->rs, mid + 1, r);
    p->update();

    return p;
}
int query(Node* p, int l, int r, int ql, int qr){
    if(!p) return -INF;
    if(ql <= l && qr >= r){
        return p->d;
    }

    int mid = l + r >> 1;
    int ans = -INF;
    if(ql <= mid) ans = query(p->ls, l, mid, ql, qr);
    if(qr > mid) ans = max(ans, query(p->rs, mid + 1, r, ql, qr));

    return ans;
}

void dfs_q(int u,int fa,int c1){
    // if(dep[u] > R) return ; // 不能加这个如果这里剪枝的话 下面的for可能就会错
    if(L <= dep[u] && dep[u] <= R) ans = max(ans, dis[u]);
    ans = max(ans, dis[u] + query(rt0, 0, n, max(L - dep[u], 0), max(0, R - dep[u])) - val[c1]);
    ans = max(ans, dis[u] + query(rt1, 0, n, max(L - dep[u], 0), max(0, R - dep[u])));

    int v, c;
    for(int i = 0;i < adj[u].size();i ++ ){
        v = adj[u][i].v, c = adj[u][i].c;
        if(v == fa || vis[v]) continue;
        dep[v] = dep[u] + 1;
        col[v] = c;
        dis[v] = dis[u] + ((col[v] != col[u]) ? val[c] : 0);
        dfs_q(v, u, c1);
    }
}

void dfs_add(int u,int fa){
    insert(rt0, 0, n, dep[u], dis[u]);
    int v;
    for(int i = 0;i < adj[u].size();i ++ ){
        v = adj[u][i].v;
        if(v == fa || vis[v]) continue;
        dfs_add(v, u);
    }
}
int last;
void solve(int u){
    // init
    top = 0;
    rt0 = rt1 = 0;
    last = -1;

    for(int i = 0;i < adj[u].size();i ++ ){
        int v = adj[u][i].v, c = adj[u][i].c;
        if(vis[v]) continue;
        if(last != -1 && c != adj[u][last].c){
            rt1 = merge(rt1, rt0, 0, n);
            rt0 = 0;
        }

        dep[v] = 1;
        col[v] = c, dis[v] = val[c];
        dfs_q(v, u, c);
        dfs_add(v, u);

        last = i;
    }
}
void divide(int u){
    // 删除当前节点
    vis[u] = true;
    solve(u);
    // cout << u << endl;
    for(int i = 0;i < adj[u].size();i ++ ){
        int v = adj[u][i].v;
        if(vis[v]) continue;
        maxp[rt = 0] = sum = siz[v];
        getrt(v, 0);
        getrt(rt, 0);
        divide(rt); // 递归分治
    }
}

int main(){
    // n 个点 m 种颜色
    scanf("%d%d%d%d", &n, &m, &L, &R);
    // 读入颜色
    for(int i = 1;i <= m;i ++ ){
        scanf("%d", &val[i]);
    }
    // 读入边  用邻接表(需要排序 用链式前向星不太好弄)
    int a, b, c;
    for(int i = 1;i < n;i ++ ){
        scanf("%d%d%d", &a, &b, &c);
        adj[a].push_back(Edge(b, c));
        adj[b].push_back(Edge(a, c));
    }
    // 排序
    for(int i = 1;i <= n;i ++ ){
        sort(adj[i].begin(), adj[i].end());
    }
    // 分治
    maxp[rt = 0] = sum = n;
    getrt(1, 0);
    getrt(rt, 0);
    divide(rt);

    printf("%d", ans);

    return 0;
}