File tree Expand file tree Collapse file tree 1 file changed +63
-0
lines changed 
Tree/3715.Sum-of-Perfect-Square-Ancestors Expand file tree Collapse file tree 1 file changed +63
-0
lines changed Original file line number Diff line number Diff line change 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+ };
    
 
   
 
     
   
   
          
     
  
    
     
 
    
      
     
 
     
    You can’t perform that action at this time.
  
 
    
  
     
    
      
        
     
 
       
      
     
   
 
    
    
  
 
  
 
     
    
0 commit comments