小红的树上路径

import sys
from collections import Counter
from itertools import combinations
from functools import lru_cache

factors_cache = {}

@lru_cache(None)
def count_factors(x):
    x = int(x)
    factors = Counter({1: 0})
    while x % 2 == 0:
        factors.update([2])
        x = x // 2

    i = 3
    while i <= x ** 0.5:
        while x % i == 0:
            factors.update([i])
            x = x // i
        i += 2
    if x != 1:
        factors.update([x])
    return (factors, count_num(factors.values()))

@lru_cache(None)
def count_num(values):
    oup = 1
    for v in values:
        oup *= v + 1
    return oup


if __name__ == '__main__':
    n, k = map(int, sys.stdin.readline().split())
    A = list(map(count_factors, sys.stdin.readline().split()))

    edges = [[] for _ in range(n)]
    for _ in range(n - 1):
        u, v = map(int, sys.stdin.readline().split())
        edges[u - 1].append(v - 1)
        edges[v - 1].append(u - 1)

    res = 0
    for mid in range(n):
        factors_mid, n_mid = A[mid]
        if n_mid >= k:
            n_neb = len(edges[mid])
            res += n_neb * (n_neb - 1) // 2
        else:
            for neb1, neb2 in combinations(edges[mid], 2):
                (factors1, n1), (factors2, n2) = A[neb1], A[neb2]
                if n1 >= k&nbs***bsp;n2 >= k&nbs***bsp;count_num((factors1 + factors_mid + factors2).values()) >= k:
                    res += 1

    print(res)

#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;
}