Minimal path on weighted tree query

Given a weighted tree with n vertices. there are q queries and for each query, you are given integers (u,k). find number of vertices v such that the smallest edge on the route from u to v is equal to k. (n,q <= 1e5)

i tried using dfs but the best solution i could think is O(n*q)

My current code:

#include <bits/stdc++.h>
 
using namespace std;

const int INF = 1e9;

struct Edge {
    int to;
    int weight;
};
 
vector<vector<Edge>> adj;
vector<int> mn;

void dfs(int u, int parent, int minWeight) {
    mn[u] = minWeight;
    for (auto edge : adj[u]) {
        if (edge.to != parent) {
            dfs(edge.to, u, min(minWeight, edge.weight));
        }
    }
}
 
int main() {
    int n, q;
    cin >> n >> q;
    adj.resize(n + 1);
    mn.resize(n + 1);
    for (int i = 0; i < n - 1; ++i) {
        int u, v, w;
        cin >> u >> v >> w;
        adj[u].push_back({v, w});
        adj[v].push_back({u, w});
    }
    while (q--) {
        int u, k;
        cin >> u >> k;
        fill(mn.begin(), mn.end(), INF);
        dfs(u, -1, INF);
        int cnt = 0;
        for (int v = 1; v <= n; ++v) {
            if (v != u && mn[v] == k) {
                cnt++;
            }
        }
 
        cout << cnt << endl;
    }
 
    return 0;
}

This can be solved offline by first reading all queries, then sorting them by edge weight in non-increasing order. We can use a disjoint set to maintain the forest formed by using only edges with weight greater than a certain value. We also sort the edges in the tree in non-increasing order and add edges of certain weights back in that order. Whenever we add edges back, we check for queries for that specific weight. The increase in component size for any node after adding these edges back is the number of vertices that have this edge weight as the minimum on the path. Note that queries for edge weights that do not exist in the tree always result in 0.

We can use a modified version of the disjoint set such that the root of each component stores the negated size of the component, to make it easier to answer queries as well as implement union by size. The time complexity of this solution is O(N log N + (N + Q) log Q + (N + Q)α(N)) (where α is the inverse Ackermann function and effectively constant here).

This can be solved online, but the code gets a lot more complicated.

#include <vector>
#include <iostream>
#include <map>
#include <functional>
#include <utility>
std::vector<int> ds; // the disjoint set
int find(int u) {
    return ds[u] < 0 ? u : ds[u] = find(ds[u]);
}
int main() {
    int n, q;
    std::cin >> n >> q;
    std::vector<int> answers(q);
    ds.assign(n + 1, -1);
    std::map<int, std::vector<std::pair<int, int>>, std::greater<>> edgesForWeight, queriesForWeight;
    for (int i = 1, u, v, w; i < n; ++i) {
        std::cin >> u >> v >> w;
        edgesForWeight[w].push_back({u, v});
    }
    for (int i = 0, u, k; i < q; ++i) {
        std::cin >> u >> k;
        queriesForWeight[k].push_back({i, u});
    }
    for (const auto& [weight, edges] : edgesForWeight) {
        auto queriesIt = queriesForWeight.find(weight);
        if (queriesIt != queriesForWeight.end())
            for (auto [qidx, node] : queriesIt->second)
                answers[qidx] = ds[find(node)];
        for (auto [u, v] : edges) {
            u = find(u), v = find(v);
            if (ds[u] > ds[v]) std::swap(u, v);
            ds[u] += ds[v];
            ds[v] = u;
        }
        if (queriesIt != queriesForWeight.end())
            for (auto [qidx, node] : queriesIt->second)
                answers[qidx] -= ds[find(node)];
    }
    for (int ans : answers) std::cout << ans << '\n';
}