Given a tree, you're asked to calculate a function $$$f(d)$$$ for $$$1\leq d \lt n$$$. Each node of the tree contains two values: $$$A_i$$$ and $$$W_i$$$. Let $$$f(d)$$$ being the sum of values of paths that have the product of $$$A$$$-values being $$$K\times t^2$$$ (with $$$t$$$ being a positive integer) and distance $$$d$$$. $$$K$$$ is given as a square-free number.

The value of a path is the product of $$$W_i$$$ for every $$$i$$$ in the path.

I've thought of $$$O(n\sqrt{n \log n})$$$, but I don't know how to optimize it more. This is done by bounding the convolution size (for FFT) in centroid decomposition. Is it really that, or is it a better time complexity if done optimally?