Skip to content

Commit 7c443de

Browse files
authored
Create 3715.Sum-of-Perfect-Square-Ancestors.cpp
1 parent 08a996f commit 7c443de

File tree

1 file changed

+63
-0
lines changed

1 file changed

+63
-0
lines changed
Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
class Solution {
2+
vector<int>freq;
3+
vector<int>sf;
4+
long long ret = 0;
5+
vector<vector<int>>next;
6+
public:
7+
void dfs(int cur, int parent){
8+
ret += freq[sf[cur]];
9+
freq[sf[cur]]+=1;
10+
for (int nxt: next[cur]) {
11+
if (nxt==parent) continue;
12+
dfs(nxt, cur);
13+
}
14+
freq[sf[cur]]--;
15+
}
16+
17+
vector<int>ComputeSpf(int N) {
18+
vector<int>spf(N+1);
19+
for (int i = 0; i <= N; ++i) spf[i] = i;
20+
for (int i = 2; i * i <= maxA; ++i) {
21+
if (spf[i] == i) {
22+
for (int j = i * i; j <= N; j += i) {
23+
if (spf[j] == j) spf[j] = i;
24+
}
25+
}
26+
}
27+
return spf;
28+
}
29+
30+
31+
long long sumOfAncestors(int n, vector<vector<int>>& edges, vector<int>& nums) {
32+
int maxA = *max_element(nums.begin(), nums.end());
33+
vector<int>spf = ComputeSpf(maxA);
34+
35+
freq.resize(maxA+1);
36+
sf.resize(n);
37+
for (int i=0; i<n; i++) {
38+
int x = nums[i];
39+
int ret = 1;
40+
while (x>1) {
41+
int p = spf[x];
42+
int cnt = 0;
43+
while (x%p==0) {
44+
x/=p;
45+
cnt++;
46+
}
47+
if (cnt%2==1) ret*=p;
48+
}
49+
sf[i] = ret;
50+
}
51+
52+
next.resize(n);
53+
for (auto& e: edges) {
54+
int u = e[0], v = e[1];
55+
next[u].push_back(v);
56+
next[v].push_back(u);
57+
}
58+
59+
dfs(0, -1);
60+
61+
return ret;
62+
}
63+
};

0 commit comments

Comments
 (0)