fork download
copy 
  1. # ========== EASY: Tree Beauty Problem ==========
  2. import sys
  3. sys.setrecursionlimit(300000)
  4. from collections import defaultdict, Counter
  5.  
  6. MOD = 10**9 + 7
  7.  
  8. def factor_remove_squares(x):
  9.     """Factor a number and remove perfect square factors"""
  10.     result = 1
  11.     d = 2
  12.     while d * d <= x:
  13.         count = 0
  14.         while x % d == 0:
  15.             count += 1
  16.             x //= d
  17.         if count % 2 == 1:
  18.             result *= d
  19.         d += 1 if d == 2 else 2
  20.     if x > 1:
  21.         result *= x
  22.     return result
  23.  
  24. def get_ans(n, par, a):
  25.     # Build tree
  26.     children = [[] for _ in range(n+1)]
  27.     for i in range(1, n):
  28.         children[par[i]].append(i+1)
  29.  
  30.     # Process node values
  31.     node_reps = [0] * (n+1)
  32.     for i in range(n):
  33.         node_reps[i+1] = factor_remove_squares(a[i])
  34.  
  35.     answer = 0
  36.  
  37.     def dfs(u):
  38.         nonlocal answer
  39.         cnt = Counter()
  40.         cnt[node_reps[u]] = 1
  41.         total_pairs = 0
  42.  
  43.         for v in children[u]:
  44.             child_cnt, child_pairs = dfs(v)
  45.             # Add child pairs to total
  46.             total_pairs += child_pairs
  47.             total_pairs %= MOD
  48.  
  49.             # Count pairs between current subtree and child subtree
  50.             for rep, count in child_cnt.items():
  51.                 if rep == 1:  # Already perfect square
  52.                     total_pairs += count
  53.                 else:
  54.                     # Find complement that makes perfect square
  55.                     complement = 1
  56.                     temp = rep
  57.                     d = 2
  58.                     while d * d <= temp:
  59.                         cnt_exp = 0
  60.                         while temp % d == 0:
  61.                             cnt_exp += 1
  62.                             temp //= d
  63.                         if cnt_exp % 2 == 1:
  64.                             complement *= d
  65.                         d += 1 if d == 2 else 2
  66.                     if temp > 1:
  67.                         complement *= temp
  68.  
  69.                     total_pairs += cnt[complement] * count
  70.  
  71.             # Merge child counter into current counter
  72.             for rep, count in child_cnt.items():
  73.                 cnt[rep] += count
  74.  
  75.         # beauty(u) = total_pairs
  76.         answer = (answer + total_pairs) % MOD
  77.         return cnt, total_pairs
  78.  
  79.     dfs(1)
  80.     return answer % MOD
Success #stdin #stdout 0.08s 14060KB
comments (?)
stdin
copy 
Standard input is empty
stdout
copy 
Standard output is empty
https://ideone.com/EQBjeC
language:
Python 3 (python  3.12)
created:
1 day ago
visibility:
public

Share or Embed source code

Discover > Sphere Engine API

The brand new service which powers Ideone!

Discover > IDE Widget

Widget for compiling and running the source code in a web browser!