不开O2,TLE了5个点,开O2,则AC,求助大佬代码能怎么优化
查看原帖
不开O2,TLE了5个点,开O2,则AC,求助大佬代码能怎么优化
822420
DGME楼主2023/8/24 02:54

如题:

#include <bits/stdc++.h>

using namespace std;

using i64 = long long;
typedef pair<int,int> pii;

constexpr int N = 250010;
vector<pii> adj[N];
int main() {
    cin.tie(0) -> sync_with_stdio(0);
    cout.tie(0) -> sync_with_stdio(0);
    int n;cin >> n;
    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});
    }
    vector<int> dfn(n + 1),dep(n + 1);
    vector<vector<int>> f(n + 1,vector<int>(20));
    vector<vector<int>> mi(n + 1,vector<int>(20,1e9));

    int idx = 0;
    function<void(int,int)> dfs = [&](int u,int fa) {
        f[u][0] = fa;
        dfn[u] = ++ idx;
        dep[u] = dep[fa] + 1;
        for(auto I : adj[u]) {
            int x = I.first;
            int v = I.second;
            if(x != fa) mi[x][0] = v,dfs(x,u);
        }
    };dfs(1,0);

    for(int i = 1;i < 20;i ++ )
        for(int j = 1;j <= n;j ++ ) {
            f[j][i] = f[f[j][i - 1]][i - 1];
            mi[j][i] = min(mi[j][i - 1],mi[f[j][i - 1]][i - 1]);
        }
    
    auto lca = [&](int l,int r) {
        if(dep[l] < dep[r]) swap(l,r);
        for(int i = 19;i >= 0;i -- ) {
            if(dep[f[l][i]] >= dep[r])
                l = f[l][i];
        }
        if(l == r) return l;
        for(int i = 19;i >= 0;i -- )
            if(f[l][i] != f[r][i]) {
                l = f[l][i];
                r = f[r][i];
            }
        return f[l][0];
    };
    
    auto cmp = [&](int a,int b) -> bool {
        return dfn[a] < dfn[b];
    };

    auto path_mi = [&](int l,int r) -> int {
        int ans = 1e9;
        for(int i = 19;i >= 0;i -- ) {
            if(dep[f[l][i]] >= dep[r]) {
                ans = min(ans,mi[l][i]);
                l = f[l][i];
            }
        }
        return ans;
    };

    vector<bool> vis(n + 1);
    vector<vector<pii>> G(n + 1);
    function<i64(int)> dfs2 = [&](int u) -> i64 {
        i64 ans = 0;
        for(auto I : G[u]) {
            int x = I.first,v = I.second;
            if(vis[x]) ans += v;
            else ans += min(dfs2(x),1LL * v);
        }
        return ans;
    };

    int m;cin >> m;
    while(m -- ) {
        int k,x;cin >> k;
        vector<int> key = {1};
        while(k -- ) {
            cin >> x,key.push_back(x);
            vis[x] = 1;
        }
        sort(key.begin(),key.end(),cmp);

        vector<int> A;
        for(int i = 0;i < key.size() - 1;i ++ ) {
            A.push_back(key[i]);
            A.push_back(lca(key[i],key[i + 1]));
        }
        A.push_back(key.back());

        sort(A.begin(),A.end(),cmp);
        A.erase(unique(A.begin(),A.end()),A.end());

        for(int i = 0;i < A.size() - 1;i ++ ) {
            int z = lca(A[i + 1],A[i]);
            int s = path_mi(A[i + 1],z);
            G[z].push_back({A[i + 1],s});
        }
        
        cout << dfs2(1) << '\n';
        for(auto x : A) G[x].clear(),vis[x] = 0;
    }
}
2023/8/24 02:54
加载中...