Codeforces Round #372 (Div. 2) E. Digit Tree （点分治）

题意

• 给一棵树，有n(n≤105)n(n≤105)个点。
• 边有权值wi(1≤wi≤9)wi(1≤wi≤9)。
• 求路径(u,v)(u,v)使得u→vu→v路径构成的大数模m为0，其中gcd（10，m）=1.
  

思路

• 点分治
• 根据欧拉定理
  xϕ(m)≡1(mod m),gcd(x,m)=1
  
  可以求出10的逆元。

#include<cmath>
#include<cstdio>
#include<cstring>
#include<set>
#include<map>
#include<queue>
#include<vector>
#include<assert.h>
#include<iostream>
#include<algorithm>
using namespace std;
#define fi first
#define se second
#define pb push_back
#define mp make_pair
#define sz(x) ((int)(x).size())
#define rep(i,l,r) for(int i=(l);i<(r);++i)
#define setIO(x) freopen(x".in","r", stdin);freopen(x".out","w",stdout);
typedef long long ll;
typedef unsigned long long ul;
typedef pair<int, int> pii;
const ul BASE = 107;
const ll LINF = 1e16 + 7;
const int N = 1e5 + 7;
const int INF = 1e9 + 7;
const int MOD = 1000000007;
const double PI = acos(-1.0);
const double EPS = 1e-8;
//-----------------head-----------------
int n, m, q[N], f[N], msz[N], siz[N], pw[N], rpw[N];
bool vis[N];
ll cnt1, cnt2;
map<int, int> R;
vector<pii> e[N];
ll kpow(ll a, ll b, ll mod) {
	ll r = 1;
	while (b > 0) {
		if (b & 1)
			r = r * a % mod;
		a = a * a % mod, b >>= 1;
	}
	return r;
}
int find(int u) {
	int t = 1;
	q[0] = u, f[u] = -1;
	rep(h, 0, t)
	{
		u = q[h];
		msz[u] = 0, siz[u] = 1;
		rep(i, 0, sz(e[u]))
		{
			int v = e[u][i].fi;
			if (vis[v] || v == f[u])
				continue;
			f[v] = u, q[t++] = v;
		}
	}
	for (int h = t - 1; h >= 0; --h) {
		u = q[h];
		msz[u] = max(msz[u], t - siz[u]);
		if (msz[u] * 2 <= t)
			return u;
		if (~f[u]) {
			siz[f[u]] += siz[u];
			msz[f[u]] = max(msz[f[u]], siz[u]);
		}
	}
	return 0;
}
void dfs1(int u, int w, int d, int p) {
	cnt1 += w == 0;
	cnt2 += R[1ll * (m - w) * rpw[d] % m];
	rep(i, 0, sz(e[u]))
	{
		int v = e[u][i].fi;
		if (vis[v] || v == p)
			continue;
		dfs1(v, (1ll * 10 * w + e[u][i].se) % m, d + 1, u);
	}
}
void dfs2(int u, int w, int d, int p) {
	cnt1 += w == 0;
	++R[w];
	rep(i, 0, sz(e[u]))
	{
		int v = e[u][i].fi;
		if (vis[v] || v == p)
			continue;
		dfs2(v, (1ll * e[u][i].se * pw[d] + w) % m, d + 1, u);
	}
}
void solve(int u) {
	u = find(u);
	R.clear();
	rep(i, 0, sz(e[u]))
	{
		int v = e[u][i].fi;
		if (vis[v])
			continue;
		dfs1(v, e[u][i].se % m, 1, u);
		dfs2(v, e[u][i].se % m, 1, u);
	}
	reverse(e[u].begin(), e[u].end());
	R.clear();
	rep(i, 0, sz(e[u]))
	{
		int v = e[u][i].fi;
		if (vis[v])
			continue;
		dfs1(v, e[u][i].se % m, 1, u);
		dfs2(v, e[u][i].se % m, 1, u);
	}
	vis[u] = true;
	rep(i, 0, sz(e[u]))
		if (!vis[e[u][i].fi])
			solve(e[u][i].fi);
}
int phi(int n) {
	int ret = n;
	for (int i = 2, I = sqrt(n + 0.5); i <= I; ++i)
		if (n % i == 0) {
			ret = ret / i * (i - 1);
			while (n % i == 0)
				n /= i;
		}
	if (n > 1)
		ret = ret / n * (n - 1);
	return ret;
}
void init(int n) {
	int inv = kpow(10, phi(m) - 1, m);
	pw[0] = rpw[0] = 1;
	rep(i, 1, n + 1)
	{
		pw[i] = 1ll * 10 * pw[i - 1] % m;
		rpw[i] = 1ll * inv * rpw[i - 1] % m;
	}
}
int main() {
	scanf("%d%d", &n, &m);
	init(n);
	rep(i, 1, n)
	{
		int u, v, w;
		scanf("%d%d%d", &u, &v, &w);
		e[u].pb(mp(v, w)), e[v].pb(mp(u, w));
	}
	memset(vis, false, sizeof(vis[0]) * n);
	cnt1 = cnt2 = 0;
	solve(0);
	printf("%I64d", cnt1 / 2 + cnt2);
	return 0;
}