AtCoder 133 F Colorful Tree

F Colorful Tree

題意： 給定一棵樹，每個樹的邊除了邊權，還有一個顏色，多條邊可能有同一種顏色，每次詢問將所有顏色爲x的邊修改爲y，詢問$u->v$的距離
分析： 樹上求距離，離不開$lca$，考慮q次修改是相互不影響的，可以離線做，先預處理出lca
然後根據慣例 $dis(u,v) = dis[u]+dis[v]-2*dis[lca(u,v)]$，怎麼處理修改呢，我們可以在$u,v,lca(u,v)$這三個點加上標記，然後進行dfs，dfs當前節點的時候我們知道$num[],value[]$ $num[i]$ 代表從1號點到當前點顏色爲i的有多少個，$value[i]$代表從1號點到現在節點,顏色$i$的貢獻是多少，之後我們就可以更新ans了

const int maxn =3e5+10;
struct Edge{
	int to,corlor,w;
};
vector<Edge> G[maxn];

int depth[maxn],vs[maxn],ID[maxn],dis[maxn];
int n;
int f[maxn][20];
void dfs(int node,int fa,int d,int &k){
   	ID[node] = k;
   	vs[k] = node;
 	depth[k++] = d;
 	// dis[node] = distance;
 	for(int i = 0;i < G[node].size(); ++i){
 		Edge &t = G[node][i];
 		if(t.to == fa) continue;
 		dis[t.to] = dis[node]+t.w;
 		dfs(t.to,node,d+1,k);
        vs[k] = node;
        depth[k++] = d;
 	}
}



void init_rmq(int n){
	
	for(int i = 1;i <= n ; ++i) f[i][0] = i;
    for(int j = 1;(1<<j) <= n; ++j){
    	for(int i = 1;i + (1<<j)-1 <= n; ++i){
    		if(depth[f[i][j-1]]< depth[f[i+(1<<(j-1))][j-1]])
    			f[i][j] = f[i][j-1];
    		else
    			f[i][j] = f[i+(1<<(j-1))][j-1];

    	}
    }
}
int rmq_query(int l,int r){
	int x = 0;
	while((1<<(x+1)) <= r-l+1) x++;
	if(depth[f[l][x]] < depth[f[r-(1<<x)+1][x]])
		return f[l][x];
	else
		return f[r-(1<<x)+1][x];
}
void init_lca(){
	int k = 1;
	dfs(1,-1,0,k);
	init_rmq(2*n-1);
}
int lca(int u,int v){
	// cout<<rmq_query(ID[u],ID[v])<<endl;
	return vs[rmq_query(min(ID[u],ID[v]),max(ID[u],ID[v]))];
}
struct que{
	int corlor,w,coe,id;
};
vector<que> Query[maxn];
int ans[maxn],num[maxn],value[maxn];// num[i] 個數，value 值
void dfs(int u,int fa){
	for(auto qu:Query[u]){
		// cout<<dis[u]<<endl;
		ans[qu.id] +=(dis[u]+num[qu.corlor]*qu.w-value[qu.corlor])*qu.coe;
	}
	for(auto p: G[u]){
		int to = p.to;
		if(to == fa) continue;
		num[p.corlor]++;
		value[p.corlor] += p.w;
		dfs(to,u);
		num[p.corlor]--;
		value[p.corlor] -= p.w;
	}
}
int main(void)
{
	int q;
    cin>>n>>q;
    for(int i = 1;i < n;++i){
    	int u,v,c,w;
    	scanf("%d%d%d%d",&u,&v,&c,&w);
    	G[u].Pb(Edge{v,c,w});
    	G[v].Pb(Edge{u,c,w});
    }
    init_lca();
    for(int i = 1;i <= q; ++i){
    	int x,y,u,v;
    	scanf("%d%d%d%d",&x,&y,&u,&v);
    	int fa = lca(u,v);

    	Query[u].Pb(que{x,y,1,i});
    	Query[v].Pb(que{x,y,1,i});
    	Query[fa].Pb(que{x,y,-2,i});
    }
    dfs(1,-1);
    for(int i  = 1;i <= q; ++i){
    	printf("%d\n",ans[i]);
    }

   return 0;
}