// @author godmar #include using namespace std; int main() { cin.tie(nullptr)->sync_with_stdio(false); long n, k; cin >> n >> k; vector V(n); for (int i = 0; i < n; i++) cin >> V[i]; vector> adj(n); for (int i = 0; i < n; i++) { int a, b; cin >> a >> b; adj[a-1].push_back(b-1); adj[b-1].push_back(a-1); } // Step 1. Find the cycle stack dfs; dfs.push(0); vector parent(n, -1); parent[0] = 0; vector::iterator> adjiter(n); for (int i = 0; i < n; i++) adjiter[i] = adj[i].begin(); vector label(n); // -1 if part of cycle, else id of "root" node unordered_map ringfreq; // frequencies in ring unordered_map treefreq; // frequencies in all trees unordered_map> val2location; // v -> { root1, root2, root3 } while (!dfs.empty()) { auto u = dfs.top(); if (adjiter[u] != adj[u].end()) { auto v = *(adjiter[u]++); if (parent[v] == -1) { parent[v] = u; dfs.push(v); } else if (parent[u] != v) { while (!dfs.empty() && dfs.top() != v) { auto w = dfs.top(); label[w] = -1; ringfreq[V[w]]++; dfs.pop(); } label[v] = -1; ringfreq[V[v]]++; break; } } else { dfs.pop(); } } // Step 2. Label each non-cycle node with the root of the tree to which it belongs. // Step 2a. Tabulate how often each height occurs vector> vfreq(n); // frequencies in each tree for (int s = 0; s < n; s++) { if (label[s] == -1) { stack> dfs; dfs.push({s, -1}); while (!dfs.empty()) { auto [u, p] = dfs.top(); dfs.pop(); if (u != s) { label[u] = s; vfreq[s][V[u]]++; treefreq[V[u]]++; val2location[V[u]].insert(s); } for (auto v : adj[u]) { if (v == p || label[v] == -1) continue; dfs.push({v, u}); } } } } // Step 3: sum up results long answer = 0; // Case 1: path starts in the ring for (auto& [v1, f1] : ringfreq) { // Case 1b: path starts and ends on the ring auto v2 = v1 + k; if (ringfreq.find(v2) != ringfreq.end()) { long comb = 2 * (long)f1 * (long)ringfreq[v2]; if (k == 0) comb -= f1; answer += comb; } // Case 1b: starts in ring, ends in tree auto _L2 = val2location.find(v2); if (_L2 != val2location.end()) { auto & L2 = _L2->second; for (auto l2 : L2) { auto f2 = vfreq[l2][v2]; long comb = 2L * (long)f1 * (long)f2; if (V[l2] == v1) comb -= f2; answer += comb; } } } // Case 2a: starts in tree and ends in tree for (auto [v1, L1] : val2location) { auto v2 = v1 + k; for (auto l1 : L1) { auto f1 = vfreq[l1][v1]; // Case 2a (1) - starts and ends in the same tree long intra = f1 * (vfreq[l1][v2]); // Case 2a (2) - starts and ends in different trees - this means it can take the ring long cross = 2L * f1 * (treefreq[v2] - vfreq[l1][v2]); answer += cross + intra; } } // Case 2b: starts in tree and ends in ring for (auto [v1, L1] : val2location) { auto v2 = v1 + k; for (auto l1 : L1) { if (ringfreq.find(v2) != ringfreq.end()) { auto f1 = vfreq[l1][v1]; auto r2 = ringfreq[v2]; long comb = 2L * (long)f1 * (long)r2; if (V[l1] == v2) // can also end in our own root, only 1 path here comb -= f1; answer += (long) comb; } } } cout << answer << endl; }