带权树最小路径边查询问题：O(nq)解法优化求助

问题描述 #

给定一棵包含n个顶点的带权树，存在q次查询，每次查询给出整数(u,k)，请找出满足u到v路径上的最小边权等于k的顶点v的数量（v≠u）。

你的当前解法使用DFS每次查询遍历整棵树，时间复杂度为O(nq)，对于n,q≤1e5的规模，1e10次操作完全超出时间限制，必须进行优化。

优化思路：离线并查集+排序处理 #

树中两点u到v的路径最小边权恰好等于k，等价于：

1. u和v在所有边权≥k的边构成的森林中连通（路径最小边权≥k）；
2. u和v在所有边权≥k+1的边构成的森林中不连通（路径最小边权<k+1，即≤k）。

因此，答案可以表示为：
答案 = (边权≥k时u所在连通块大小 - 1) - (边权≥k+1时u所在连通块大小 - 1) = 边权≥k的连通块大小 - 边权≥k+1的连通块大小

我们采用离线处理来高效计算所有查询：

1. 将所有边按权值从大到小排序；
2. 收集所有查询的k值及其k+1值，去重后按降序排序；
3. 使用带大小维护的并查集，按权值从大到小依次合并连通块；
4. 对每个关键k值，记录查询涉及节点的连通块大小；
5. 最后计算每个查询的答案。

优化后的代码实现 #

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

const int MAXN = 1e5 + 5;

struct Edge {
    int u, v, w;
    bool operator<(const Edge& other) const {
        return w > other.w; // 边权降序排序
    }
};

struct Query {
    int u, k, idx;
};

int parent[MAXN], sz[MAXN];
Edge edges[MAXN];
Query queries[MAXN];
int ans[MAXN];

// 并查集查找（路径压缩）
int find(int x) {
    return parent[x] == x ? x : (parent[x] = find(parent[x]));
}

// 并查集合并（按大小合并）
void unite(int x, int y) {
    x = find(x), y = find(y);
    if (x == y) return;
    if (sz[x] < sz[y]) swap(x, y);
    parent[y] = x;
    sz[x] += sz[y];
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int n, q;
    cin >> n >> q;

    // 读入边
    for (int i = 0; i < n-1; ++i) {
        cin >> edges[i].u >> edges[i].v >> edges[i].w;
    }
    sort(edges, edges + n-1);

    // 读入查询
    for (int i = 0; i < q; ++i) {
        cin >> queries[i].u >> queries[i].k;
        queries[i].idx = i;
    }

    // 收集所有需要处理的关键k值：查询的k和k+1
    set<int> key_ks;
    for (auto& qry : queries) {
        key_ks.insert(qry.k);
        key_ks.insert(qry.k + 1);
    }
    vector<int> sorted_ks(key_ks.begin(), key_ks.end());
    sort(sorted_ks.rbegin(), sorted_ks.rend()); // 降序排序

    // 存储每个关键k对应的查询节点的连通块大小
    unordered_map<int, unordered_map<int, int>> k_to_size;

    // 初始化并查集
    for (int i = 1; i <= n; ++i) {
        parent[i] = i;
        sz[i] = 1;
    }

    int edge_ptr = 0;
    for (int k : sorted_ks) {
        // 合并所有边权≥当前k的边
        while (edge_ptr < n-1 && edges[edge_ptr].w >= k) {
            unite(edges[edge_ptr].u, edges[edge_ptr].v);
            edge_ptr++;
        }

        // 只记录当前k相关的查询节点的大小，节省空间
        unordered_set<int> needed_nodes;
        for (auto& qry : queries) {
            if (qry.k == k || qry.k + 1 == k) {
                needed_nodes.insert(qry.u);
            }
        }
        for (int u : needed_nodes) {
            k_to_size[k][u] = sz[find(u)];
        }
    }

    // 计算每个查询的答案
    for (auto& qry : queries) {
        int u = qry.u;
        int k = qry.k;
        int size_k = k_to_size[k][u];
        int size_k1 = k_to_size[k+1][u];
        ans[qry.idx] = size_k - size_k1;
    }

    // 输出答案
    for (int i = 0; i < q; ++i) {
        cout << ans[i] << '\n';
    }

    return 0;
}

复杂度分析 #

• 边排序：O(n log n)；
• 关键k值排序：O(q log q)；
• 并查集操作：O(n α(n))，α为阿克曼函数，近似常数；
• 总时间复杂度：O(n log n + q log q)，完全满足n,q≤1e5的时间要求。

内容的提问来源于stack exchange，提问作者turtle silver