CF715C Digit Tree

传送门

题目大意

给你一棵带边权的树，求有多少点对 $(u,v)$ 满足从 $u$ 到 $v$ 路径上经过的边依次写出来构成的十进制数能够被给出的整数 $M$ 整除。

Solution

从 dsu on tree 过来的，所以直接开始想 dsu 的暴力。

首先很显然我们会想到用两个数组来维护依次写出十进制数，设 $u{p}_x$ 表示从 $x$ 到根节点所构成数，$dow{n}_x$ 表示从根节点到 $x$ 所构成的数。两者可以简单一次大法师解决。具体可以看代码，重点不是这个，所以就不展开讲。

接下来我们考虑，如果一个点对 $(st,ed)$ 是满足要求的，那么用上面的两个数组来表示就是（$u$ 表示 $st$ 和 $ed$ 的 $LCA$）：
$$\frac{u{p}_{st}-u{p}_u}{10^{de{p}_u}}\times 10^{de{p}_{ed}-de{p}_u}+dow{n}_{ed}-dow{n}_u\times 10^{de{p}_{ed}-de{p}_u}\equiv 0(\mathrm{m}\mathrm{o}\mathrm{d}M)$$

比较复杂，但是解释一下就很简单了。

我们考虑把路径分为 $st\to u$ 和 $u\to ed$，那么 $st\to ed$ 路径上构成的十进制数就是 $(st\to u)\times 10^{de{p}_{ed}-de{p}_u}+(u\to ed)$，这很好理解吧。

然后分别求出 $st\to u$ 和 $u\to ed$ 的十进制数。这也很简单，分别是 $\frac{u{p}_{st}-u{p}_u}{10^{de{p}_u}}$ 和 $dow{n}_{ed}-dow{n}_u\times 10^{de{p}_{ed}-de{p}_u}$。画个图就完了。

然后考虑 dsu 的时候，我们需要对于一个点，快速查询有多少点能与之形成合法的路径，那我们分两种情况讨论：

1. 如果当前节点作为 $st$，那么疯狂移项得到：$$dow{n}_{ed}\times 10^{-de{p}_{ed}}\equiv (dow{n}_u-\frac{u{p}_{st}-u{p}_u}{10^{de{p}_u}})\times 10^{-de{p}_u}(\mathrm{m}\mathrm{o}\mathrm{d}M)$$
   所以我们对于前面的点计数 $dow{n}_x\times 10^{-de{p}_x}$ 加一，之后查询后面一坨东西在计数器中有多少和它相等的（当然先要模上 $M$）。

2. 否则作为 $ed$，那么疯狂移项有：$$u{p}_{st}\equiv dow{n}_u\times 10^{de{p}_u}-dow{n}_{ed}\times 10^{2\times de{p}_u-de{p}_{ed}}+u{p}_u(\mathrm{m}\mathrm{o}\mathrm{d}M)$$
   与上同理，我们塞进一个计数器里。

所以，我们用 $map$ 作为计数器就可以了。

为了节约时间，不用每次求 $LCA$，我们直接钦定当前子树的根节点作为 $LCA$ 就可以了。这样我们在处理子树的时候，先统计答案，再加上子树的信息，可以保证每次匹配的两个不在一个子树中，也就是 $LCA$ 为当前根。而如果在同一子树中，在前面 dsu 的时候已经计算过了。

Code

#include<bits/stdc++.h>
#define ll long long
#define inf (1<<30)
#define INF (1ll<<60)
using namespace std;
const int MAXN=1e5+10;
ll MOD;
ll EXGCD(ll a,ll b,ll &x,ll &y){
	if(b==0){x=1;y=0;return a;}
	ll r=EXGCD(b,a%b,y,x);
	y-=a/b*x;return r;
}
ll inv(ll x){
	ll a,k;ll r=EXGCD(x,MOD,a,k);
	a=(a%MOD+MOD)%MOD;
	return a;
}
vector<pair<int,int> > e[MAXN];
ll up[MAXN],down[MAXN],tp[MAXN];
int siz[MAXN],son[MAXN],dep[MAXN];
void dfs(int x,int fa){
	siz[x]=1;son[x]=-1;
	for(auto s:e[x]){
		if(s.first==fa) continue;
		dep[s.first]=dep[x]+1;
		up[s.first]=(up[x]+tp[dep[x]]*s.second%MOD)%MOD;
		down[s.first]=(down[x]*10%MOD+s.second)%MOD;
		dfs(s.first,x);
		siz[x]+=siz[s.first];
		if(son[x]==-1||siz[son[x]]<siz[s.first])
			son[x]=s.first;
	}
}
ll p10(int x){
	if(x>=0) return tp[x];
	else return inv(tp[-x]);
}
map<long long,int> S,T;
ll ans=0;
void calc(int x,int u){
	ll st=((down[u]-(((up[x]-up[u])%MOD+MOD)%MOD)*inv(tp[dep[u]])%MOD)%MOD+MOD)%MOD*inv(tp[dep[u]])%MOD;
	ll ed=(((down[u]*tp[dep[u]]%MOD+up[u])%MOD-down[x]*p10(2*dep[u]-dep[x])%MOD)%MOD+MOD)%MOD;
	ans+=S[ed];ans+=T[st];
}
void getans(int x,int fa,int u){
	calc(x,u);
	for(auto s:e[x])
		if(s.first!=fa)
			getans(s.first,x,u);
}
void add(int x,int fa){
	S[up[x]]++;
	T[down[x]*inv(tp[dep[x]])%MOD]++;
	for(auto s:e[x])
		if(s.first!=fa)
			add(s.first,x);
}
void sub(int x,int fa){
	S[up[x]]--;
	T[down[x]*inv(tp[dep[x]])%MOD]--;
	for(auto s:e[x])
		if(s.first!=fa)
			sub(s.first,x);
}
void dsu(int x,int fa){
	for(auto s:e[x]){
		if(s.first==fa) continue;
		if(s.first==son[x]) continue;
		dsu(s.first,x);sub(s.first,x);
	}if(~son[x]) dsu(son[x],x);
	for(auto s:e[x]){
		if(s.first==fa) continue;
		if(s.first==son[x]) continue;
		getans(s.first,x,x);
		add(s.first,x);
	}calc(x,x);
	S[up[x]]++;
	T[down[x]*inv(tp[dep[x]])%MOD]++;
}
int main()
{
	int n;
	scanf("%d%lld",&n,&MOD);
	tp[0]=1;
	for(int i=1;i<=n;i++)
		tp[i]=tp[i-1]*10%MOD;
	for(int i=1,u,v,w;i<n;i++){
		scanf("%d%d%d",&u,&v,&w);u++;v++;
		e[u].push_back(make_pair(v,w));
		e[v].push_back(make_pair(u,w));
	}dfs(1,0);dsu(1,0);
	printf("%lld\n",ans);
}