Path Queries

Created by LXC on Sat Mar 30 00:12:06 2024

https://codeforces.com/problemset/problem/1213/G

ranting: 1800

tag: divide and conquer, dsu, graphs, sortings, trees

problem

$\mathsf E \color{red}\mathsf{ntropyIncreaser}$ 有一棵 $n$ 个点的树，每条边都带权。

她会问你 $m$ 个问题，每次给你一个正整数 $q$，求最大权值不大于 $q$ 的简单路径数量。

需要注意的是，对于一个点对 $(u,v)$ 只记一次，单独一个点不算路径。

输入格式

第一行两个正整数 $n,m$，意义如题目描述。
接下来 $n-1$ 行，每行三个正整数 $u,v,w$，表示 $u,v$ 之间有一条权为 $w$ 的无向边。
最后一行 $m$ 个正整数，表示询问。

输出格式

对于每个询问，输出一行一个整数表示答案。

数据范围

$1\le n,m \le 2\times10^5$
$1\le u,v \le n$
$1\le w,q \le 2\times 10^5$

solution

离线操作，对于查询排序，对于查询x，连接边权小于等于x的所有边，计算各个联通分量的贡献，对于一个联通块中含k个点，那么就贡献了$\binom{k}{2}$条路径。

code

#include <bits/stdc++.h>
#define SINGLE_INPUT
#define ll long long
#define ull unsigned long long
#define N 500005
#define MOD 998244353
using namespace std;

template<class t,class u> ostream& operator<<(ostream& os,const pair<t,u>& p) {
    return os<<'['<<p.first<<", "<<p.second<<']';
}
template<class t> ostream& operator<<(ostream& os,const vector<t>& v) {
    os<<'['; int s = 1;
    for(auto e:v) { if (s) s = 0; else os << ", "; os << e; }
    return os<<']';
}
template<class t,class u> ostream& operator<<(ostream& os,const map<t,u>& mp){
    os<<'{'; int s = 1;
    for(auto [x,y]:mp) { if (s) s = 0; else os << ", "; os<<x<<": "<<y; }
    return os<<'}';
}

struct Edge {
    int x, y, w;
    Edge(int x, int y, int w):x(x),y(y),w(w){}
    bool operator<(const Edge& o) const {
        return w < o.w;
    }
};

void sol() {
    int n, q;
    cin >> n >> q;
    vector<int> fa(n+1, -1);
    function<int(int)> ufind = [&](int x) {
        return fa[x] < 0 ? x : fa[x] = ufind(fa[x]);
    };
    vector<Edge> edge;
    for (int i=1; i<n; i++) {
        int x, y, w;
        cin >> x >> y >> w;
        edge.emplace_back(x, y, w);
    }
    sort(edge.begin(), edge.end());
    vector<ll> qr(q), idx(q), ans(q, 0);
    for (auto& i:qr) cin >> i;
    iota(idx.begin(), idx.end(), 0);
    sort(idx.begin(), idx.end(), [&](ll a, ll b) {
        return qr[a] < qr[b];
    });
    auto cpt = [&](ll x)->ll {
        return x*(x-1)/2;
    };
    // cout << qr << endl;
    // cout << idx << endl;
    ll cur = 0;
    int p = 0;
    for (int i:idx) {
        while (p<n-1 && edge[p].w <= qr[i]) {
            int x = edge[p].x, y = edge[p].y;
            int fx = ufind(x), fy = ufind(y);
            cur -= cpt(-fa[fx]);
            cur -= cpt(-fa[fy]);
            fa[fx] += fa[fy];
            fa[fy] = fx;
            cur += cpt(-fa[fx]);
            p++;
        }
        ans[i] = cur;
    }
    for (auto i:ans) {
        cout << i << " ";
    }
    cout << "\n";
}

int main() {
    cout << setprecision(15) << fixed;
    ios::sync_with_stdio(false);
    cin.tie(0);
    cout.tie(0);
#ifndef SINGLE_INPUT
    int t;
    cin >> t;
    while (t--) {
        sol();
    }
#else
    sol();
#endif
    return 0;
}