Virtual Tree (Auxiliary Tree)

1
#include <bits/stdc++.h>
2
using namespace std;
3
4
/*
5
  Virtual Tree (Auxiliary Tree) construction on an unweighted tree.
6
  - Preprocess LCA with binary lifting.
7
  - Given k marked nodes, build the virtual tree in O(k log k).
8
  - Compute the total number of original edges in the minimal subtree connecting them
9
    (Steiner subtree length) as sum of depth differences across VT edges.
10
*/
11
12
struct LCA {
13
    int n, LOG, root;
14
    vector<vector<int>> up;   // up[v][j] = 2^j-th ancestor of v
15
    vector<int> depth, tin, tout;
16
    vector<vector<int>> g;
17
    int timer = 0;
18
19
    LCA(int n=0): n(n) {}
20
21
    void init(int _n, int _root, const vector<vector<int>>& _g) {
22
        n = _n; root = _root; g = _g;
23
        LOG = 1; while ((1<<LOG) <= n) ++LOG;
24
        up.assign(n+1, vector<int>(LOG));
25
        depth.assign(n+1, 0);
26
        tin.assign(n+1, 0);
27
        tout.assign(n+1, 0);
28
        timer = 0;
29
        dfs(root, root);
30
    }
31
32
    void dfs(int v, int p) {
33
        tin[v] = ++timer;
34
        up[v][0] = p;
35
        for (int j = 1; j < LOG; ++j) up[v][j] = up[ up[v][j-1] ][j-1];
36
        for (int to : g[v]) if (to != p) {
37
            depth[to] = depth[v] + 1; // unweighted: each edge has length 1
38
            dfs(to, v);
39
        }
40
        tout[v] = timer;
41
    }
42
43
    bool is_ancestor(int u, int v) const { // u is ancestor of v?
44
        return tin[u] <= tin[v] && tout[v] <= tout[u];
45
    }
46
47
    int lca(int u, int v) const {
48
        if (is_ancestor(u, v)) return u;
49
        if (is_ancestor(v, u)) return v;
50
        int x = u;
51
        for (int j = LOG-1; j >= 0; --j) {
52
            int a = up[x][j];
53
            if (!is_ancestor(a, v)) x = a;
54
        }
55
        return up[x][0];
56
    }
57
};
58
59
int main() {
60
    ios::sync_with_stdio(false);
61
    cin.tie(nullptr);
62
63
    int n; cin >> n;
64
    vector<vector<int>> g(n+1);
65
    for (int i = 0; i < n-1; ++i) {
66
        int u, v; cin >> u >> v;
67
        g[u].push_back(v);
68
        g[v].push_back(u);
69
    }
70
71
    int root = 1;
72
    LCA solver; solver.init(n, root, g);
73
74
    int k; cin >> k;
75
    vector<int> a(k);
76
    for (int i = 0; i < k; ++i) cin >> a[i];
77
78
    // Step 1: sort by tin and deduplicate
79
    sort(a.begin(), a.end(), [&](int x, int y){ return solver.tin[x] < solver.tin[y]; });
80
    a.erase(unique(a.begin(), a.end()), a.end());
81
82
    // Step 2: insert LCAs of consecutive nodes
83
    vector<int> nodes = a;
84
    for (int i = 0; i+1 < (int)a.size(); ++i) nodes.push_back(solver.lca(a[i], a[i+1]));
85
    sort(nodes.begin(), nodes.end(), [&](int x, int y){ return solver.tin[x] < solver.tin[y]; });
86
    nodes.erase(unique(nodes.begin(), nodes.end()), nodes.end());
87
88
    // Step 3: build Virtual Tree adjacency (directed parent->child)
89
    auto is_anc = [&](int u, int v){ return solver.is_ancestor(u, v); };
90
    vector<vector<int>> vt(n+1); // allocate big but we'll only touch `nodes`
91
92
    vector<int> st;
93
    st.reserve(nodes.size());
94
    st.push_back(nodes[0]);
95
    for (int i = 1; i < (int)nodes.size(); ++i) {
96
        int v = nodes[i];
97
        while ((int)st.size() >= 1 && !is_anc(st.back(), v)) {
98
            int u = st.back(); st.pop_back();
99
            // connect u to its closest ancestor now at st.back()
100
            int p = st.back();
101
            vt[p].push_back(u);
102
        }
103
        st.push_back(v);
104
    }
105
    while ((int)st.size() > 1) {
106
        int u = st.back(); st.pop_back();
107
        int p = st.back();
108
        vt[p].push_back(u);
109
    }
110
111
    // Step 4: Compute Steiner subtree length = sum over edges depth[v]-depth[p]
112
    long long steiner_len = 0;
113
    for (int p : nodes) {
114
        for (int v : vt[p]) {
115
            steiner_len += (long long)(solver.depth[v] - solver.depth[p]);
116
        }
117
    }
118
119
    cout << steiner_len << "\n";
120
    return 0;
121
}
122

1
#include <bits/stdc++.h>
2
using namespace std;
3
4
/*
5
  Compute sum of pairwise distances among k marked nodes.
6
  Key identity on Virtual Tree:
7
    Sum_{pairs (x,y)} dist(x,y) = sum over VT edges (p->v): s_v * (k - s_v) * len(p,v)
8
  where s_v is the number of marked nodes in v's component below the edge (p,v),
9
  and len(p,v) = depth[v] - depth[p] (unweighted).
10
*/
11
12
struct LCA {
13
    int n, LOG, root, timer = 0;
14
    vector<vector<int>> up, g;
15
    vector<int> tin, tout, depth;
16
    void init(int _n, int _root, const vector<vector<int>>& _g){
17
        n=_n; root=_root; g=_g;
18
        LOG=1; while((1<<LOG)<=n) ++LOG;
19
        up.assign(n+1, vector<int>(LOG));
20
        tin.assign(n+1,0); tout.assign(n+1,0); depth.assign(n+1,0);
21
        dfs(root, root);
22
    }
23
    void dfs(int v, int p){
24
        tin[v]=++timer; up[v][0]=p;
25
        for(int j=1;j<LOG;++j) up[v][j]=up[ up[v][j-1] ][j-1];
26
        for(int to: g[v]) if(to!=p){ depth[to]=depth[v]+1; dfs(to,v);} 
27
        tout[v]=timer;
28
    }
29
    bool is_ancestor(int u, int v) const { return tin[u]<=tin[v] && tout[v]<=tout[u]; }
30
    int lca(int u,int v) const{
31
        if(is_ancestor(u,v)) return u; if(is_ancestor(v,u)) return v;
32
        int x=u; for(int j=LOG-1;j>=0;--j){int a=up[x][j]; if(!is_ancestor(a,v)) x=a;} return up[x][0];
33
    }
34
};
35
36
int main(){
37
    ios::sync_with_stdio(false);
38
    cin.tie(nullptr);
39
40
    int n; cin>>n; 
41
    vector<vector<int>> g(n+1);
42
    for(int i=0;i<n-1;++i){int u,v;cin>>u>>v; g[u].push_back(v); g[v].push_back(u);} 
43
44
    LCA lca; lca.init(n, 1, g);
45
46
    int k; cin>>k; vector<int> a(k); for(int i=0;i<k;++i) cin>>a[i];
47
48
    // sort & unique by tin
49
    sort(a.begin(), a.end(), [&](int x,int y){return lca.tin[x]<lca.tin[y];});
50
    a.erase(unique(a.begin(), a.end()), a.end());
51
52
    // insert LCAs of consecutive
53
    vector<int> nodes=a; 
54
    for(int i=0;i+1<(int)a.size();++i) nodes.push_back(lca.lca(a[i], a[i+1]));
55
    sort(nodes.begin(), nodes.end(), [&](int x,int y){return lca.tin[x]<lca.tin[y];});
56
    nodes.erase(unique(nodes.begin(), nodes.end()), nodes.end());
57
58
    // build VT
59
    vector<vector<int>> vt(n+1);
60
    vector<int> st; st.push_back(nodes[0]);
61
    auto is_anc = [&](int u,int v){return lca.is_ancestor(u,v);} ;
62
    for(int i=1;i<(int)nodes.size();++i){
63
        int v=nodes[i];
64
        while(!is_anc(st.back(), v)){
65
            int u=st.back(); st.pop_back();
66
            int p=st.back(); vt[p].push_back(u);
67
        }
68
        st.push_back(v);
69
    }
70
    while((int)st.size()>1){int u=st.back(); st.pop_back(); int p=st.back(); vt[p].push_back(u);}    
71
72
    // mark important nodes
73
    unordered_set<int> S(a.begin(), a.end());
74
75
    // DFS over VT to compute s_v and accumulate answer
76
    long long ans=0; 
77
    function<int(int)> dfs = [&](int u){
78
        int cnt = S.count(u) ? 1 : 0;
79
        for(int v: vt[u]){
80
            int sub = dfs(v);
81
            long long len = (long long)(lca.depth[v]-lca.depth[u]);
82
            ans += len * (long long)sub * (long long)( (int)a.size() - sub );
83
            cnt += sub;
84
        }
85
        return cnt;
86
    };
87
88
    dfs(nodes[0]);
89
    cout<<ans<<"\n";
90
    return 0;
91
}
92

1
#include <bits/stdc++.h>
2
using namespace std;
3
4
/*
5
  Supports multiple queries on a weighted tree using a Virtual Tree per query.
6
  - Preprocess LCA with distances-to-root D[v] (weighted depths).
7
  - For each query's marked set S:
8
      * Build VT.
9
      * Compute Steiner subtree total weight.
10
      * Compute sum of pairwise weighted distances among S.
11
  Efficient clearing: only touch nodes present in VT.
12
*/
13
14
struct LCA {
15
    int n, LOG, root, timer = 0;
16
    vector<vector<int>> up;
17
    vector<int> tin, tout, depth;
18
    vector<long long> D; // distance to root (weighted)
19
    vector<vector<pair<int,long long>>> g;
20
21
    void init(int _n, int _root, const vector<vector<pair<int,long long>>>& _g){
22
        n=_n; root=_root; g=_g;
23
        LOG=1; while((1<<LOG)<=n) ++LOG;
24
        up.assign(n+1, vector<int>(LOG));
25
        tin.assign(n+1,0); tout.assign(n+1,0); depth.assign(n+1,0);
26
        D.assign(n+1,0);
27
        dfs(root, root);
28
    }
29
    void dfs(int v,int p){
30
        tin[v]=++timer; up[v][0]=p;
31
        for(int j=1;j<LOG;++j) up[v][j]=up[ up[v][j-1] ][j-1];
32
        for(auto [to,w]: g[v]) if(to!=p){
33
            depth[to]=depth[v]+1; D[to]=D[v]+w; dfs(to,v);
34
        }
35
        tout[v]=timer;
36
    }
37
    bool is_ancestor(int u,int v) const {return tin[u]<=tin[v] && tout[v]<=tout[u];}
38
    int lca(int u,int v) const{
39
        if(is_ancestor(u,v)) return u; if(is_ancestor(v,u)) return v;
40
        int x=u; for(int j=LOG-1;j>=0;--j){int a=up[x][j]; if(!is_ancestor(a,v)) x=a;} return up[x][0];
41
    }
42
    long long dist(int u,int v) const { return D[u]+D[v]-2LL*D[lca(u,v)]; }
43
};
44
45
int main(){
46
    ios::sync_with_stdio(false);
47
    cin.tie(nullptr);
48
49
    int n; cin>>n;
50
    vector<vector<pair<int,long long>>> g(n+1);
51
    for(int i=0;i<n-1;++i){
52
        int u,v; long long w; cin>>u>>v>>w;
53
        g[u].push_back({v,w}); g[v].push_back({u,w});
54
    }
55
    LCA lca; lca.init(n, 1, g);
56
57
    int Q; cin>>Q;
58
59
    // We'll reuse these per query to avoid reallocation overhead
60
    vector<vector<pair<int,long long>>> vt(n+1);
61
    vector<int> touched; touched.reserve(2e5);
62
63
    while(Q--){
64
        int k; cin>>k; vector<int> a(k);
65
        for(int i=0;i<k;++i) cin>>a[i];
66
        sort(a.begin(), a.end(), [&](int x,int y){return lca.tin[x]<lca.tin[y];});
67
        a.erase(unique(a.begin(), a.end()), a.end());
68
69
        vector<int> nodes=a;
70
        for(int i=0;i+1<(int)a.size();++i) nodes.push_back(lca.lca(a[i], a[i+1]));
71
        sort(nodes.begin(), nodes.end(), [&](int x,int y){return lca.tin[x]<lca.tin[y];});
72
        nodes.erase(unique(nodes.begin(), nodes.end()), nodes.end());
73
74
        // Clean only touched from previous query
75
        for(int v: touched) vt[v].clear();
76
        touched.clear();
77
        for(int v: nodes) touched.push_back(v);
78
79
        // Build weighted VT: edge weight is dist(parent, child) in original tree
80
        vector<int> st; st.push_back(nodes[0]);
81
        auto is_anc = [&](int u,int v){return lca.is_ancestor(u,v);} ;
82
        for(int i=1;i<(int)nodes.size();++i){
83
            int v=nodes[i];
84
            while(!is_anc(st.back(), v)){
85
                int u=st.back(); st.pop_back();
86
                int p=st.back(); long long w = lca.dist(p,u);
87
                vt[p].push_back({u,w});
88
            }
89
            st.push_back(v);
90
        }
91
        while((int)st.size()>1){int u=st.back(); st.pop_back(); int p=st.back(); long long w=lca.dist(p,u); vt[p].push_back({u,w});}
92
93
        // Marked set lookup
94
        unordered_set<int> S(a.begin(), a.end());
95
        int K = (int)S.size();
96
97
        // Compute Steiner length (weighted) and pairwise sum (weighted)
98
        long long steiner_len = 0;
99
        long long pair_sum = 0;
100
101
        function<int(int)> dfs = [&](int u){
102
            int cnt = S.count(u) ? 1 : 0;
103
            for(auto [v,w]: vt[u]){
104
                int sub = dfs(v);
105
                steiner_len += w;                      // each VT edge counted once
106
                pair_sum += w * 1LL * sub * 1LL * (K - sub); // contribution formula
107
                cnt += sub;
108
            }
109
            return cnt;
110
        };
111
        dfs(nodes[0]);
112
113
        cout << steiner_len << " " << pair_sum << "\n";
114
    }
115
    return 0;
116
}
117