PROBLEM LINK:
Contest
Practice
Setter: iceknight1093,, jainmilind
Testers: tabr
Editorialist: aryanag_adm
DIFFICULTY:
2782
PREREQUISITES:
Factorisation, Small to Large
PROBLEM:
You are given a rooted tree T with N vertices. Each vertex v has an integer A_v associated with it. For each v, you need to output the number of distinct divisors of \prod_{u \in \text{subtree of }v} A_u.
EXPLANATION:
We can precompute the prime factorisation of all numbers between 1 and 10^6 in 10^6 \times \log 10^6 steps, which is fast enough. We will assume we have these prime factorisations available.
Now, if x = \prod_{i=1}^k p_i^{y_i} where p_i are prime and y_i are positive integers, then the number of distinct divisors of x is \prod_{i=1}^k (y_i + 1).
Now, let’s say we are given the prime factorisation (modulo 10^9 + 7) and number of divisors d of a large number X. We want to compute the prime factorisation and number of divisors d' of X \cdot x, where x \leq 10^6. Let’s say X = \prod_{p \in P} p^{y_p}, x = \prod_{p \in P'} p^{z_p}. Then d' = d \cdot (\prod_{p \in P' \setminus P} {z_p + 1}) \cdot (\prod_{p \in P \cap P'} \frac{x_p + y_p + 1}{y_p + 1}). The complete prime factorisation can be similarly computed. Note that there are at most P' algebraic operations here, and P' \leq \log 10^6 < 20. Therefore, we can do this fast.
Now our algorithm is easy. If v is a leaf, the answer can be computed directly using prime factorisation of A_v. For any non leaf node v, we can use small-to-large along with the merging algorithm we saw above to compute the answer fast.
TIME COMPLEXITY:
Roughly, the time complexity is O(1e6 \log 1e6 + 20 N \log^3 N ).
SOLUTION:
Editorialist's Solution
#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
//#include <sys/resource.h>
#define int long long
#define double long double
#define initrand mt19937 mt_rand(time(0));
#define rand mt_rand()
#define MOD 1000000007
#define INF 1000000000
#define mid(l, u) ((l+u)/2)
#define rchild(i) (i*2 + 2)
#define lchild(i) (i*2 + 1)
#define mp(a, b) make_pair(a, b)
#define lz lazup(l, u, i);
#define ordered_set tree<pair<int, int>, null_type,less<pair<int, int>>, rb_tree_tag,tree_order_statistics_node_update>
using namespace std;
using namespace __gnu_pbds;
const int maxn = 3e4 + 5;
int ans[maxn], A[maxn], sz[maxn], st[maxn], ft[maxn], ansd[(int)1e6 + 1], tt = 0, currdiv = 1;;
vector<int> adj[maxn], ver;
vector<pair<int, int>> divisors[(int)1e6 + 1];
bool notPrime[(int)1e6 + 1];
map<int, int> currfac;
int ex(int a, int b){
if(b==0) return 1;
int c = ex(a, b/2);
c *= c;
c %= MOD;
if(b%2) c *= a;
c %= MOD;
return c;
}
int inverse(int x){
return ex(x, MOD - 2);
}
int d(int a, int b){
return (a*inverse(b))%MOD;
}
int m(int a, int b){
return (a*b)%MOD;
}
void getsz(int v, int p){
st[v] = tt++;
ver.push_back(v);
sz[v] = 1; // every vertex has itself in its subtree
for(auto u : adj[v])
if(u != p){
getsz(u, v);
sz[v] += sz[u]; // add size of child u to its parent(v)
}
ft[v] = tt;
}
void dfs(int v, int p, bool keep){
int mx = -1, bigChild = -1;
for(auto u : adj[v])
if(u != p && sz[u] > mx)
mx = sz[u], bigChild = u;
for(auto u : adj[v])
if(u != p && u != bigChild)
dfs(u, v, 0); // run a dfs on small childs and clear them from cnt
if(bigChild != -1)
dfs(bigChild, v, 1); // bigChild marked as big and not cleared from cnt
for(auto u : adj[v])
if(u != p && u != bigChild)
for(int p = st[u]; p < ft[u]; p++)
for(pair<int, int> j: divisors[A[ver[p]]]){
currdiv = d(currdiv, currfac[j.first] + 1);
currfac[j.first] += j.second;
currdiv = m(currdiv, currfac[j.first] + 1);
}
for(pair<int, int> j: divisors[A[v]]){
currdiv = d(currdiv, currfac[j.first] + 1);
currfac[j.first] += j.second;
currdiv = m(currdiv, currfac[j.first] + 1);
}
ans[v] = currdiv;
//now cnt[c] is the number of vertices in subtree of vertex v that has color c. You can answer the queries easily.
if(keep == 0) {
currdiv = 1;
currfac.clear();
}
}
signed main(){
ios_base::sync_with_stdio(false);
cin.tie(NULL);
cout.tie(NULL);
for(int i = 1;i<=1e6;i++) ansd[i] = 1;
for(int i = 2;i<=1e6;i++){
if(notPrime[i]) continue;
for(int j = i;j<=1e6;j+=i){
notPrime[j] = true;
int mult = 0;
int temp = j;
while((temp%i) == 0){
temp /= i;
mult++;
}
ansd[j] *= (mult + 1);
ansd[j] %= MOD;
divisors[j].push_back({i, mult});
}
}
int t;
cin>>t;
while(t--){
int n;
cin>>n;
ver.clear();
currfac.clear();
tt = 0;
for(int i = 0;i<n;i++) adj[i].clear();
for(int i = 0;i<n;i++) cin>>A[i];
for(int i = 1;i<n;i++){
int u, v;
cin>>u>>v;
u--;
v--;
adj[u].push_back(v);
adj[v].push_back(u);
}
tt = 0;
getsz(0, 0);
currdiv = 1;
dfs(0, 0, 9);
for(int i = 0;i<n;i++){
cout<<ans[i]<<" ";
}
cout<<endl;
}
}