3715. Sum of Perfect Square Ancestors

You are given an integer n and an undirected tree rooted at node 0 with n nodes numbered from 0 to n - 1. This is represented by a 2D array edges of length n - 1, where edges[i] = [ui, vi] indicates an undirected edge between nodes ui and vi.

You are also given an integer array nums, where nums[i] is the positive integer assigned to node i.

Define a value ti as the number of ancestors of node i such that the product nums[i] * nums[ancestor] is a perfect square.

Return the sum of all ti values for all nodes i in range [1, n - 1].

Note:

• In a rooted tree, the ancestors of node i are all nodes on the path from node i to the root node 0, excluding i itself.

class Solution {
    long long dfs(vector<vector<int>>& adj, unordered_map<int,int>& freq, vector<int>& nums, int u, int par) {
        long long res = freq[nums[u]];
        freq[nums[u]]++;
        for(auto& v : adj[u]) {
            if(v == par) continue;
            res += dfs(adj,freq,nums,v,u);
        }
        freq[nums[u]]--;
        return res;
    }
public:
    long long sumOfAncestors(int n, vector<vector<int>>& edges, vector<int>& nums) {
        vector<int> pows;
        for(int i = 2; i * i <= 1e5; i++) {
            pows.push_back(i * i);
        }
        for(auto& n : nums) {
            for(auto& p : pows) {
                if(n < p) break;
                while(n % p == 0) n /= p;
            }
        }
        vector<vector<int>> adj(n);
        for(auto& e : edges) {
            int u = e[0], v = e[1];
            adj[u].push_back(v);
            adj[v].push_back(u);
        }
        unordered_map<int,int> freq;
        return dfs(adj,freq,nums,0,-1);
    }
};