2019 南昌邀请赛 E Interesting Trip

题意

给定一棵 $n$ 个点的无根树,每个点有点权,求长度为 $D$ 且满足路径点权的 $\gcd$ 大于 $1$ 的路径个数。

其中 $3 \le n \le 5 \cdot 10^5$,值域 $[1,30000]$。

分析

记 $f(n)$ 表示路径 $\gcd$ 和为 $n$ 的路径数,$g(n)$ 表示路径 $\gcd$ 为 $n$ 的倍数的路径数。

因此有

$$
g(d)=\sum_{d|n} f(d)
$$

由莫比乌斯反演

$$
f(d)=\sum_{d|n} \mu(\frac{n}{d}) g(n)
$$

因此答案 $S$ 就是

$$
S=\sum_{i=2}^\infty f(i)=\sum_{i=2}^\infty \sum_{i|d} \mu(\frac{d}{i}) g(d)
$$

交换求和顺序

$$
S=\sum_{d=1}^\infty g(d) \sum_{i>1 \land i|d} \mu(\frac{d}{i}) = \sum_{d=2}^\infty g(d)(0-\mu(d))
$$

于是只需要计算 $g(d)$ 即可,枚举值域 $i \in [2,30000]$ 且 $\mu(i) \neq 0$,计算有多少条长度为 $D$ 的路径的点权都是 $i$ 的倍数。将这些点拿出来,显然是这个点集构成一个森林,下面考虑一棵树的情况。记 $dp(u,i)$ 表示以 $u$ 为根节点的子树内距离 $u$ 为 $i$ 的点数。状态数和转移的复杂度都是 $O(n^2)$ 的,套用长链剖分优化关于深度信息合并的技巧,可以做到 $O(n)$ 的时空求出这个 $dp$ 数组。

赛中乱写了个点分治冲过去了,计蒜客机子卡常卡到死,靠评测姬抖了抖终于过去了。

代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
#include <iostream>
#include <cstdio>
#include <cstring>
#include <cmath>
#include <algorithm>
#include <vector>
#include <string>
#include <utility>
#include <functional>
#define ms(a,b) memset(a,b,sizeof(a))
using namespace std;
typedef long long ll;
typedef pair<int,int> PII;
const int mod = 998244353;
const int inf = 1 << 30;
const int maxn = 500000 + 5;
const int maxm = 30000 + 5;

struct fastIO {
char s[100000]; int it,len;
fastIO() { it = len = 0; }
inline char get() {
if (it < len) return s[it++];
it = 0; len = fread(s, 1, 100000, stdin);
if (len == 0) return EOF;
else return s[it++];
}
bool notend() {
char c = get();
while(c == ' ' || c == '\n') c = get();
if (it > 0) it--;
return c != EOF;
}
} buff;
inline int gi() {
int r = 0; bool ng = 0;
char c = buff.get();
while (c != '-' && (c < '0' || c > '9')) c = buff.get();
if (c == '-') ng = 1, c = buff.get();
while (c >= '0' && c <= '9') r = r * 10 + c - '0', c = buff.get();
return ng ? -r : r;
}

namespace Mu {
int mu[maxm], vis[maxm], prime[maxm], tot;
vector<int> mud[maxm];
void getMu() {
tot = 0; mu[1] = 1;
for (int i = 2; i < maxm; ++i) {
if (!vis[i]) prime[++tot] = i, mu[i] = -1;
for (int j = 1; j <= tot && i * prime[j] < maxm; j++) {
vis[i * prime[j]] = 1;
if (i % prime[j] == 0) {
mu[i * prime[j]] = 0; break;
}
mu[i * prime[j]] = -mu[i];
}
}
for (int i = 1; i < maxm; i++) {
for (int x = 1; x * x <= i; x++) {
if (i % x) continue;
if (mu[x]) mud[i].push_back(x);
if (x * x != i && mu[i / x]) mud[i].push_back(i / x);
}
}
}
}
using Mu::mu;
using Mu::mud;

int n, D, a[maxn];
int vis[maxn], mxd[maxn], son[maxn], dp[maxn];
int* f[maxn];

struct Edge {
int to, nxt;
} e[maxn * 2];
int ecnt, head[maxn];
void adde(int u, int v) {
e[++ecnt] = { v, head[u] };
head[u] = ecnt;
}
vector<int> ig[maxm];

int main() {
Mu::getMu();
int T = gi(), kase = 0;
while (T--) {
n = gi(); D = gi();
ecnt = 0;
for (int i = 1; i <= n; i++) head[i] = vis[i] = 0;
for (int i = 1; i <= 30000; i++) ig[i].clear();
for (int i = 1; i <= n; i++) {
a[i] = gi();
for (auto x: mud[a[i]]) ig[x].push_back(i);
}
for (int i = 2, u, v; i <= n; i++) {
u = gi(); v = gi();
adde(u, v); adde(v, u);
}
ll ans = 0;
for (int i = 2; i <= 30000; i++) {
if (ig[i].empty()) continue;
int xs = -mu[i];
for (auto& x: ig[i]) {
if (vis[x] == i) continue;
int cnt = 0;
function<void(int,int)> predfs = [&](int u, int ff) {
cnt++; vis[u] = i; son[u] = mxd[u] = 0;
for (int j = head[u]; j; j = e[j].nxt) {
int v = e[j].to;
if (v == ff || a[v] % i) continue;
predfs(v, u);
if (mxd[v] > mxd[u]) {
son[u] = v; mxd[u] = mxd[v];
}
}
mxd[u]++;
};
int *tot = dp;
auto make = [&](int u) {
f[u] = tot; tot += mxd[u];
};
function<void(int,int)> dfs = [&](int u, int ff) {
if (son[u]) {
int v = son[u];
f[v] = f[u] + 1;
dfs(v, u);
}
if (mxd[u] - 1 >= D) ans += xs * f[u][D];
f[u][0] = 1;
for (int j = head[u]; j; j = e[j].nxt) {
int v = e[j].to;
if (v == ff || v == son[u] || a[v] % i) continue;
make(v); dfs(v, u);
for (int i = max(D - mxd[u], 0); i < mxd[v] && i < D; i++) {
ans += 1ll * xs * f[v][i] * f[u][D - i - 1];
}
for (int i = 0; i < mxd[v] && i < D; i++) {
f[u][i + 1] += f[v][i];
}
}
};
predfs(x, 0);
make(x); for (int i = 0; i <= cnt; i++) dp[i] = 0;
dfs(x, 0);
}
}
printf("Case #%d: %lld\n", ++kase, 2ll * ans);
}
return 0;
}