CF715C Digit Tree

题面

题解

树上求点对数目的题大多都是点分治解决

路径统计中有一个问题，如果现在求出从上到下的链长$a$，从下到上的链长$b$和深度$c$，

则：

$ a \times10^c + b\equiv0(mod\;m) $

两边同除以$c$，得

$ a + b\times10^{-c}\equiv0(mod\;m)\;\;(\because gcd(10,\;m) = 1) $

于是可以很好维护了。

代码

#include<cstdio>
#include<cstring>
#include<cctype>
#include<algorithm>
#include<map>
#define RG register
#define file(x) freopen(#x".in", "r", stdin);freopen(#x".out", "w", stdout);
#define clear(x, y) memset(x, y, sizeof(x))

inline int read()
{
	int data = 0, w = 1; char ch = getchar();
	while(ch != '-' && (!isdigit(ch))) ch = getchar();
	if(ch == '-') w = -1, ch = getchar();
	while(isdigit(ch)) data = data * 10 + (ch ^ 48), ch = getchar();
	return data * w;
}

const int maxn(1e5 + 10);
struct edge { int next, to, dis; } e[maxn << 1];
int head[maxn], e_num, n, m, size[maxn], root;
int Size, min, vis[maxn], Pow[maxn], cnt;

inline void Init()
{
	clear(head, 0); e_num = 0, Pow[0] = 1;
	for(RG int i = 1; i <= n; i++) Pow[i] = 10ll * Pow[i - 1] % m;
}

#undef clear

inline void add_edge(int from, int to, int dis)
{
	e[++e_num] = (edge) {head[from], to, dis};
	head[from] = e_num;
}

std::pair<int, int> digit[maxn << 1];
std::map<int, long long> tong;
void getRoot(int x, int fa)
{
	size[x] = 1;
	int tot = 0;
	for(RG int i = head[x]; i; i = e[i].next)
	{
		int to = e[i].to;
		if(to == fa || vis[to]) continue;
		getRoot(to, x); tot = std::max(tot, size[to]);
		size[x] += size[to];
	}
	tot = std::max(tot, Size - size[x]);
	if(tot < min) min = tot, root = x;
}

void getDep(int x, int fa, long long d1, long long d2, int dep)
{
	if(dep >= 0) ++tong[d1], digit[++cnt] = std::make_pair(d2, dep);
	for(RG int i = head[x]; i; i = e[i].next)
	{
		int to = e[i].to, dis = e[i].dis; if(to == fa || vis[to]) continue;
		getDep(to, x, (d1 + 1ll * Pow[dep + 1] * dis) % m,
				(d2 * 10 + dis) % m, dep + 1);
	}
}

void exgcd(long long a, long long b, long long &gcd, long long &x, long long &y)
{
	!b ? gcd = a, x = 1, y = 0 : (exgcd(b, a % b, gcd, y, x), y -= x * (a / b));
}

inline long long Inv(int a, int Mod)
{
	long long d, x, y; exgcd(a, Mod, d, x, y);
	return d == 1 ? (x % Mod + Mod) % Mod : -1;
}

long long calc(int x, int d)
{
	long long ans = 0; tong.clear(), cnt = 0;
	if(d) getDep(x, 0, d % m, d % m, 0);
	else getDep(x, 0, 0, 0, -1);
	for(RG int i = 1; i <= cnt; i++)
	{
		long long tmp = (-digit[i].first * 1ll *
		Inv(Pow[digit[i].second + 1], m) % m + m) % m;
		if(tong.find(tmp) != tong.end()) ans += tong[tmp];
		if(!d) ans += (digit[i].first == 0);
	}
	if(d == 0) ans += tong[0];
	return ans;
}

long long ans;
void dfs(int x)
{
	ans += calc(x, 0);
	vis[x] = true;
	for(RG int i = head[x]; i; i = e[i].next)
	{
		int to = e[i].to;
		if(vis[to]) continue;
		ans -= calc(to, e[i].dis);
		min = n, Size = size[to];
		getRoot(to, x); dfs(root);
	}
}

int main()
{
	while(~scanf("%d%d", &n, &m))
	{
		Init(); Size = min = n;
		for(RG int i = 1, a, b, c; i < n; i++)
			a = read() + 1, b = read() + 1, c = read(),
			  add_edge(a, b, c), add_edge(b, a, c);
		ans = 0; memset(vis, 0, sizeof(vis));
		getRoot(1, 0); dfs(root);
		printf("%I64d\n", ans);
	}
	return 0;
}