1778F - Maximizing Root - CodeForces Solution

#include <bits/stdc++.h>

using namespace std;

const int INF = 1e9 + 5;

int t, n, k, rt, divcnt, ans;

vector<int> arr;
vector<vector<int>> divs, tree, dp;

void dfs(int v, int p){
    if(tree[v].size() == 1 && v != 1){
        for(int i = 0; i < divcnt; i++){
            if(arr[v] % divs[rt][i] == 0) dp[v][i] = 0;
        }
    }
    else{
        fill(dp[v].begin(), dp[v].end(), 0);

        for(auto e : tree[v]){
            if(e != p){
                dfs(e, v);

                for(int i = 0; i < divcnt; i++){
                    if(arr[v] % divs[rt][i] != 0) dp[v][i] = -1;

                    if(dp[v][i] != -1){
                        int mn = INF;

                        for(int j = i; j < divcnt; j++){
                            if(divs[rt][j] % divs[rt][i] == 0 && dp[e][j] != -1){
                                mn = min(mn, dp[e][j]);
                            }
                        }

                        if(mn == INF) dp[v][i] = -1;
                        else dp[v][i] += mn;
                    }
                }
            }
        }
    }

    if(v != 1){
        vector<int> newdp(divcnt, INF);

        for(int i = 0; i < divcnt; i++){
            if(dp[v][i] != -1){
                for(int j = 0; j < divcnt; j++){
                    if((divs[rt][i] * divs[rt][i]) % divs[rt][j] == 0){
                        newdp[j] = min(newdp[j], dp[v][i] + 1);
                    }
                }
            }
        }

        for(int i = 0; i < divcnt; i++){
            if(newdp[i] != INF){
                if(dp[v][i] == -1) dp[v][i] = newdp[i];
                else dp[v][i] = min(dp[v][i], newdp[i]);
            }
        }
    }
}

void solve(){
    cin >> n >> k;

    arr.assign(n + 1, 0);
    tree.assign(n + 1, {});

    for(int i = 1; i <= n; i++) cin >> arr[i];

    rt = arr[1];
    divcnt = divs[rt].size();

    for(int i = 0, a, b; i < n - 1; i++){
        cin >> a >> b;

        tree[a].push_back(b);
        tree[b].push_back(a);
    }

    dp.assign(n + 1, vector(divcnt, -1));

    dfs(1, 0);

    ans = 0;

    for(int i = divcnt - 1; i >= 0; i--){
        if(dp[1][i] != -1){
            if(dp[1][i] <= k - 1) ans = max(ans, rt * divs[rt][i]);
            else if(dp[1][i] <= k) ans = max(ans, rt);
        }
    }

    cout << ans << "\n";
}

int main(){
    ios_base::sync_with_stdio(false);cin.tie(NULL);

    divs.assign(1000 + 1, {});

    for(int i = 1; i <= 1000; i++){
        for(int j = 1; j <= i; j++){
            if(i % j == 0) divs[i].push_back(j);
        }
    }

    cin >> t;

    while(t--){
        solve();
    }
}