小红的树上路径

#include <iostream>
#include <vector>
#include <cstring>
#include <utility>
using namespace std;

const int MAXN = 1000001;
int div_cnt[MAXN];

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

    int n, k;
    cin >> n >> k;
    vector<int> a(n+1);
    for (int i=1; i<=n; i++) {
        cin >> a[i];
    }
    vector<vector<int>> graph(n+1);
    for (int i=0; i<n-1; i++) {
        int u, v;
        cin >> u >> v;
        graph[u].push_back(v);
        graph[v].push_back(u);
    }

    memset(div_cnt, 0, sizeof(div_cnt));
    for (int i=1; i<MAXN; i++) {
        for (int j=i; j<MAXN; j+=i) {
            div_cnt[j]++;
        }
    }

    vector<pair<int, int>> lists[101];
    for (int w=1; w<=100; w++) {
        for (int i=1; i<=100; i++) {
            for (int j=i; j<=100; j++) {
                int product = w * i * j;
                if (product < MAXN && div_cnt[product] >= k) {
                    lists[w].emplace_back(i, j);
                }
            }
        }
    }

    long long ans = 0;
    for (int v=1; v<=n; v++) {
        int w_val = a[v];
        int cnt[101] = {0};
        for (int u : graph[v]) {
            int w_u = a[u];
            if (w_u >= 1 && w_u <= 100) {
                cnt[w_u]++;
            }
        }
        long long S_v = 0;
        for (auto &p : lists[w_val]) {
            int i = p.first, j = p.second;
            if (i == j) {
                S_v += (long long) cnt[i] * (cnt[i] - 1) / 2;
            } else {
                S_v += (long long) cnt[i] * cnt[j];
            }
        }
        ans += S_v;
    }

    cout << ans << endl;
    return 0;
}