Submission #1294712

#TimeUsernameProblemLanguageResultExecution timeMemory
1294712harryleeeFactories (JOI14_factories)C++20
100 / 100
6853 ms217736 KiB
#include <bits/stdc++.h>
#include "factories.h"
using namespace std;
const int maxn = 500000 + 5;
const int LOG = 20;
const long long INFLL = (long long)1e18;

int n;
vector<pair<int,long long>> adj[maxn];

int h[maxn];
long long dis[maxn];
int st[maxn][LOG];

// centroid decomp
bool delNode[maxn];
int subSize[maxn];
int parCent[maxn]; // parent in centroid-tree for centroid nodes (optional)
vector<int> centroidsOfNode[maxn];
long long bestAtCent[maxn];

// ---------- LCA ----------
void dfs_lca(int root){
    // iterative stack to avoid recursion overflow
    stack<pair<int,int>> s;
    s.push({root, -1});
    h[root] = 0;
    dis[root] = 0;
    st[root][0] = root;
    while(!s.empty()){
        auto [u, p] = s.top(); s.pop();
        for (auto &e : adj[u]){
            int v = e.first;
            long long w = e.second;
            if (v == p) continue;
            h[v] = h[u] + 1;
            dis[v] = dis[u] + w;
            st[v][0] = u;
            s.push({v, u});
        }
    }
    // fill binary lifting table
    for (int j = 1; j < LOG; ++j){
        for (int i = 0; i < n; ++i){
            st[i][j] = st[ st[i][j-1] ][j-1];
        }
    }
}

int lca(int u, int v){
    if (h[u] > h[v]) swap(u,v);
    int k = h[v] - h[u];
    for (int i = 0; i < LOG; ++i) if (k & (1<<i)) v = st[v][i];
    if (u == v) return u;
    for (int i = LOG-1; i >= 0; --i){
        if (st[u][i] != st[v][i]){
            u = st[u][i];
            v = st[v][i];
        }
    }
    return st[u][0];
}

inline long long length_uv(int u, int v){
    int w = lca(u,v);
    return dis[u] + dis[v] - 2*dis[w];
}

// ---------- centroid decomposition ----------
void dfs_size(int u, int p){
    subSize[u] = 1;
    for (auto &e : adj[u]){
        int v = e.first;
        if (v == p || delNode[v]) continue;
        dfs_size(v, u);
        subSize[u] += subSize[v];
    }
}

int find_centroid(int u, int p, int total){
    for (auto &e : adj[u]){
        int v = e.first;
        if (v == p || delNode[v]) continue;
        if (subSize[v] > total/2) return find_centroid(v, u, total);
    }
    return u;
}

void mark_component_with_centroid(int u, int p, int centroid){
    // add centroid into all nodes of this component
    centroidsOfNode[u].push_back(centroid);
    for (auto &e : adj[u]){
        int v = e.first;
        if (v == p || delNode[v]) continue;
        mark_component_with_centroid(v, u, centroid);
    }
}

void decompose(int entry, int parentC){
    dfs_size(entry, -1);
    int c = find_centroid(entry, -1, subSize[entry]);
    parCent[c] = parentC;
    // mark all nodes in this component with centroid c
    mark_component_with_centroid(c, -1, c);
    // now remove centroid and recurse on components
    delNode[c] = true;
    for (auto &e : adj[c]){
        int v = e.first;
        if (!delNode[v]){
            decompose(v, c);
        }
    }
    // optionally delNode[c] stays true
}

// ---------- interface ----------
void Init(int N, int A[], int B[], int D[]){
    n = N;
    // clear structures
    for (int i = 0; i < n; ++i){
        adj[i].clear();
        centroidsOfNode[i].clear();
        delNode[i] = false;
        bestAtCent[i] = INFLL;
        for (int j = 0; j < LOG; ++j) st[i][j] = i; // temp init
    }
    // build edges (2-way)
    for (int i = 0; i < n-1; ++i){
        int u = A[i], v = B[i];
        long long w = D[i];
        adj[u].push_back({v,w});
        adj[v].push_back({u,w});
    }
    // LCA preprocess (root at 0). Use iterative DFS to set st, h, dis
    dfs_lca(0);
    // centroid decomposition
    decompose(0, -1);
    // bestAtCent already set to INF
}

long long Query(int S, int X[], int T, int Y[]){
    vector<int> touched;
    touched.reserve((S+T) * 20);
    // update centroids with Y nodes
    for (int i = 0; i < T; ++i){
        int y = Y[i];
        for (int c : centroidsOfNode[y]){
            long long val = length_uv(y, c);
            if (bestAtCent[c] == INFLL) touched.push_back(c);
            if (val < bestAtCent[c]) bestAtCent[c] = val;
        }
    }
    long long res = INFLL;
    for (int i = 0; i < S; ++i){
        int x = X[i];
        for (int c : centroidsOfNode[x]){
            if (bestAtCent[c] == INFLL) continue;
            long long total = bestAtCent[c] + length_uv(x, c);
            if (total < res) res = total;
        }
    }
    // reset touched centroids
    for (int c : touched) bestAtCent[c] = INFLL;
    return res;
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...