【计蒜之道2019初赛1 BCD】【计蒜客39263】商汤AI园区的n个路口 题解

题目大意

  有一棵 $n$ 个点的树，每条边有边权，边权互不相同，范围为 $[1,m]$。
  现在你要给每个点定一个点权，点权范围也是 $[1,m]$。
  假设一条边连着 $a$ 和 $b$，边权为 $w$，那么点权 $v_a$ 和 $v_b$ 要满足 $\gcd(v_a,v_b)\not = w$。
  求方案数。

  medium：
  $1 \leq n \leq m \leq 1000$
  hard：
  $1 \leq n \leq m \leq 10^5$

$$\\$$ $$\\$$ $$\\$$

  这个题挺有趣的。

medium

  首先列出一个 dp：设 $f_{i,j}$ 表示第 $i$ 个点，点权为 $j$ 的方案数，记 $s_i=\sum_{j=1}^m f_{i,j}$，那么转移很简单：

$$\begin{aligned} f_{i,j}&=\prod_{son}\sum_{k=1}^m [\gcd(j,k)\not=w]f_{son,k} \\ &=\prod_{son}(s_{son}-\sum_{k=1}^m [\gcd(j,k)=w]f_{son,k}) \end{aligned}$$

  计算减去 $[\gcd(j,k)=w]f_{son,k}$ 这一部分的时候，实际上 $j$ 和 $k$ 都是 $w$ 的倍数。由于边权互不相同，因此时间最多为

$$\begin{aligned} &\sum_{w=1}^m \lfloor \frac mw \rfloor^2 \\ \leq &\bigg(\sum_{w=1}^m \lfloor \frac mw \rfloor \bigg)^2 \\ \leq &(m\ln m)^2 \end{aligned}$$

  因此总的时间为 $O(nm+m^2\ln^2 m)$，这就可以过掉 medium 了。

hard

  考虑优化这个 dp。可以发现，大部分的 $f_{i,j}$ 都等于 $\prod_{son}s_{son}$，只有 $m \ln m$ 个 $f_{i,j}$ 是特殊的。所以只要快速计算出这些特殊位置就行了。
  即，对于 $i$ 的儿子 $son$（连过去的边的边权为 $w$），当 $j$ 为 $w$ 的倍数时，快速得到

$$illegal_j=\sum_{k=1}^m [\gcd(j,k)=w]f_{son,k}$$

  莫比乌斯反演一下：

$$\begin{aligned} &\sum_{k=1}^m [\gcd(j,k)=w]f_{son,k} \\ =& \sum_{k=1}^{\lfloor \frac mw \rfloor} [\gcd(\frac jw,k)=1] f_{son,kw} \\ =& \sum_{d|\frac jw} \mu(d) \sum_{k=1}^{\lfloor \frac{m}{wd} \rfloor} f_{son,kwd} \end{aligned}$$

  对于 $son$，枚举 $d$，然后枚举 $k$ 算出 $\mu(d) \sum_{k=1}^{\lfloor \frac{m}{wd} \rfloor} f_{son,kwd}$，最后枚举 $j$ 把这个 $d$ 的贡献加到那个 $j$ 里就好了。

  来算一下复杂度。$d$ 的范围是 $[1,\lfloor \frac mw \rfloor]$，在枚举了 $d$ 以后枚举 $k$ 的时间就是

$$\sum_{d=1}^{\lfloor \frac mw \rfloor} \lfloor \frac{m}{wd} \rfloor=O(\frac mw \ln \frac mw)$$

  枚举 $j$ 同理，算上全部的 $w$ 的话就是

$$\begin{aligned} &\sum_{w=1}^m O(\frac mw \ln \frac mw) \\ =& O(\ln m) \cdot \sum_{w=1}^m O(\frac mw) \\ =& O(\ln m) \cdot O(m \ln m) \\ =& O(m \ln^2 m) \end{aligned}$$

  因此总时间是 $O(n+m\ln^2m)$

代码

#include<bits/stdc++.h>
#define fo(i,a,b) for(int i=a;i<=b;i++)
using namespace std;

typedef long long LL;

const int maxn=1e5+5;
const LL mo=1e9+7;

int n,m;

int p0,p[maxn],mu[maxn];
bool bz[maxn];
void Pre(int n)
{
	mu[1]=1;
	fo(i,2,n)
	{
		if (!bz[i]) p[++p0]=i, mu[i]=-1;
		fo(j,1,p0)
		{
			if ((LL)i*p[j]>n) break;
			bz[i*p[j]]=1;
			if (i%p[j]==0) break;
				else mu[i*p[j]]=-mu[i];
		}
	}
}

int tot,go[2*maxn],val[2*maxn],nxt[2*maxn],f1[maxn];
void ins(int x,int y,int z)
{
	go[++tot]=y;
	val[tot]=z;
	nxt[tot]=f1[x];
	f1[x]=tot;
}

LL mi(LL x,LL y)
{
	LL re=1;
	for(; y; y>>=1, x=x*x%mo) if (y&1) re=re*x%mo;
	return re;
}

map<int,LL> f[maxn];
LL s[maxn],ilg[maxn];
LL findF(int k,int j) { return (f[k].count(j)) ?f[k][j] :f[k][0] ; }
void dfs(int k,int last)
{
	LL pro=1;
	for(int p=f1[k]; p; p=nxt[p]) if (go[p]!=last)
	{
		dfs(go[p],k);
		
		(pro*=s[go[p]])%=mo;
		for(int d=val[p], di=1; d<=m; d+=val[p], di++)
		{
			LL sum=0;
			for(int y=d; y<=m; y+=d) (sum+=mu[di]*findF(go[p],y))%=mo;
			for(int x=d; x<=m; x+=d) (ilg[x/val[p]]+=sum)%=mo;
		}
		LL ny=mi(s[go[p]],mo-2);
		for(int j=val[p]; j<=m; j+=val[p])
			if (f[k].count(j)) (f[k][j]*=(s[go[p]]-ilg[j/val[p]])*ny%mo)%=mo;
				else f[k][j]=(s[go[p]]-ilg[j/val[p]])*ny%mo;
		fo(j,1,m/val[p]) ilg[j]=0;
	}
	s[k]=pro*m%mo;
	for(auto p:f[k])
	{
		int j=p.first;
		(f[k][j]*=pro)%=mo;
		s[k]=(s[k]-pro+f[k][j])%mo;
	}
	f[k][0]=pro;
}

int main()
{
	Pre(maxn-5);
	
	scanf("%d %d",&n,&m);
	fo(i,1,n-1)
	{
		int x,y,w;
		scanf("%d %d %d",&x,&y,&w);
		ins(x,y,w), ins(y,x,w);
	}
	
	dfs(1,0);
	
	printf("%lld\n",(s[1]+mo)%mo);
}