CF715C Digit Tree

$\text{Link}$

看到题解区的大佬都用 $\text{map}$，蒟蒻表示只会 $\text{sort}$，于是这里给出一个不用 $\text{map}$ 的点分治做法。

前置知识

• 树的重心。
• 点分治。
• 乘法逆元。

题意

给定一棵 $n$ 个点的树，任意边权 $w$ 满足 $1 \leq w\leq 9$。给定一个数 $m$，保证 $\gcd(m,10)=1$。

对于一对不同的点对 $(u,v)$（$(u,v)$ 和 $(v,u)$ 也是不同的），将 $u$ 到 $v$ 的最短路 $u\Rightarrow v$ 经过的边的边权连接起来，得到一个数字 $x$，若满足 $x\equiv0\pmod m$，则称 $(u,v)$ 是有趣的点对。

求这颗树上有多少对有趣的点对？

$n≤10^5,m\leq 10^9$。

分析

统计树上满足要求的点对个数常常要用点分治，这题也不例外。

提前预处理 $10$ 的幂及逆元。

接下来我们考虑如何计算一棵树中经过重心 $center$ 的点对个数。

首先，预处理出每个点的层数 $d_i$ 和重心到第 $i$ 个点返往分别形成的数 $num1_i,num2_i$，此时，$(u,v)$ 的 $x$ 等于 $num2_u\times 10^{d_v}+num1_v$，即 $num2_u\times 10^{d_v}+num1_v\equiv0\pmod m$，因为 $\gcd(m,10)=1$，所以 $num2_u\equiv-num1_v\times 10^{-d_v}\pmod m$，枚举每个 $num1_v$，找出相应的 $num2_u$ 的个数即可。

如何求呢？在预处理层数的同时，我们记录每个节点属于 $center$ 的哪个子节点节点，也就是 $i$ 的颜色 $c_i$，我们把每个 $num2_i$ 存入数组 $a_{c_i}$ 中（$a_{c_i}$ 表示以 $center$ 的子节点 $c_i$ 为根的子树的 $num2_i$ 集合），存入数组 $a2$ 中（$a2$ 表示的是以 $center$ 为根的树中的 $num2_i$ 集合），同时把每个 $\text{pair}$ 对 $(num1_i,i)$ 存入数组 $a1$ 中（$a1$ 表示以 $center$ 为根的树中的 $num1_i$ 集合）。

接下来把 $a_i,a1,a2$ 数组排序，枚举 $a1$ 里的元素 $(num1_i,i)$ ，此时我们要找的数 $need\equiv-num1_i\times 10^{-d_i}\pmod m$，我们需要的是子树 $c_i$ 外的 $need$ 的个数，根据容斥原理，子树 $c_i$ 外的 $need$ 的个数等于以 $center$ 为根的树减去在子树 $c_i$ 内 $need$ 的个数，这两个部分可以通过 $a_i,a1,a2$ 在二分查找得到。注意，如果某个 $num1_i$ 或 $num2_i$ 的是 $m$ 的倍数，说明 $i\Rightarrow center$ 或 $i\Leftarrow center$ 这条链的端点本身就是有趣的点对，直接让答案增加 $1$。

代码

#include<bits/stdc++.h>
#define ll long long
using namespace std;
long long read(){
	long long x=0,f=1;char ch=getchar();
	while(!isdigit(ch)){if(ch=='-') f=-1;ch=getchar();}
	while(isdigit(ch)){x=x*10+ch-48;ch=getchar();}
	return x*f;
}
void write(long long x){
    if(x<0) putchar('-'),x=-x;
    if(x>9) write(x/10);
    putchar(x%10+'0');
}
const int N=1e5+10;
int n,s[N],t,c[N];
ll m,ans,d[N],num1[N],num2[N],m10[N],nim10[N];
vector<int>e[N],a[N],a2;
vector<ll>w[N];
vector<pair<ll,int> >a1;
pair<ll,int> gg;
bool v1[N],v2[N];
long long exgcd(long long a,long long b,long long &x,long long &y){
    if(!b){
        x=1;y=0;
        return a;
    }
    long long d=exgcd(b,a%b,x,y);
    long long z=x;x=y;y=z-y*(a/b);
    return d;
}
long long qni_ol(long long a,long long p){
    long long x,y;
    exgcd(a,p,x,y);
    return x;
}
void add(int u,int v,ll W){
	e[u].push_back(v);
	w[u].push_back(W);
}
struct Tree_Center{
	int sum;
	int size[N],weight[N];
	int center;
	void dfs(int x,int fa){
		size[x]=1;
		weight[x]=0;
		for(int i=0;i<e[x].size();i++){
			if(e[x][i]==fa||v2[e[x][i]])
				continue;
			dfs(e[x][i],x);
			size[x]+=size[e[x][i]];
			weight[x]=max(weight[x],size[e[x][i]]);
		}
		weight[x]=max(weight[x],sum-size[x]);
		if(weight[x]<weight[center])
			center=x;
	}
	void ask(int root,int n){
		sum=n;
		weight[center=0]=1e9;
		dfs(root,0);
	}
}tree;
void dfs(int x,int fa,int col){
	s[x]=1;
	d[x]=d[fa]+1;
	for(int i=0;i<e[x].size();i++){
		if(e[x][i]==fa||v2[e[x][i]])
			continue;
		num1[e[x][i]]=(num1[x]*10%m+w[x][i])%m;
		num2[e[x][i]]=(num2[x]+w[x][i]*m10[d[x]]%m)%m;
		c[e[x][i]]=col;
		a[col].push_back(num2[e[x][i]]);
		gg.first=num1[e[x][i]];gg.second=e[x][i];
		a1.push_back(gg);
		a2.push_back(num2[e[x][i]]);
		dfs(e[x][i],x,col);
		s[x]+=s[e[x][i]];
	}
}
ll get(int col,ll need){
	return upper_bound(a[col].begin(),a[col].end(),need)-
	lower_bound(a[col].begin(),a[col].end(),need);
}
void calc(int x){
	d[x]=0;
	while(a1.size())
		a1.pop_back();
	while(a2.size())
		a2.pop_back();
	for(int i=0;i<e[x].size();i++){
		if(v2[e[x][i]])
			continue;
		while(a[e[x][i]].size())
			a[e[x][i]].pop_back();
		num1[e[x][i]]=num2[e[x][i]]=w[x][i]%m;
		c[e[x][i]]=e[x][i];
		a[e[x][i]].push_back(num2[e[x][i]]);
		gg.first=num1[e[x][i]];gg.second=e[x][i];
		a1.push_back(gg);
		a2.push_back(num2[e[x][i]]);
		dfs(e[x][i],x,e[x][i]);
		sort(a[e[x][i]].begin(),a[e[x][i]].end());
	}
	sort(a1.begin(),a1.end());
	sort(a2.begin(),a2.end());
	ll need;
	for(int i=0;i<a1.size();i++){
		if(!a1[i].first)
			ans++;
		need=(m-a1[i].first*nim10[d[a1[i].second]]%m)%m;
		ans+=upper_bound(a2.begin(),a2.end(),need)-
		lower_bound(a2.begin(),a2.end(),need)-get(c[a1[i].second],need);
	}
	for(int i=0;i<a2.size();i++)
		if(!a2[i])
			ans++;
}
void DFZ(int x,int sum){
	tree.ask(x,sum);
	int center=tree.center;
	calc(center);
	v2[center]=1;
	for(int i=0;i<e[center].size();i++){
		if(v2[e[center][i]])
			continue;
		DFZ(e[center][i],s[e[center][i]]);
	}
}
int main(){
	n=read();m=read();
	m10[0]=1;
	for(int i=1;i<=n;i++){
		m10[i]=m10[i-1]*10%m;
		nim10[i]=qni_ol(m10[i],m);
	}
	int u,v,w;
	for(int i=1;i<n;i++){
		u=read()+1;v=read()+1;w=read();
		add(u,v,w);add(v,u,w);
	}
	DFZ(1,n);
	write(ans);
	return 0;
}