Submission #1230893

#	Time	Username	Problem	Language	Result	Execution time	Memory
1230893		wcarrotw	Factories (JOI14_factories)	C++17	100 / 100	2555 ms	161900 KiB

factories

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

#define ll long long
const int N = 5e5 + 5;
const int LG = 20;
const ll INF = 1e18;

int n, timer;
vector<pair<int, ll>> g[N];
int tin[N], tout[N], depth[N], parent[N][LG];
ll dist[N];

void dfs(int u, int p) {
    tin[u] = ++timer;
    parent[u][0] = p;
    for (int i = 1; i < LG; ++i)
        parent[u][i] = parent[parent[u][i - 1]][i - 1];
    for (auto [v, w] : g[u]) {
        if (v == p) continue;
        depth[v] = depth[u] + 1;
        dist[v] = dist[u] + w;
        dfs(v, u);
    }
    tout[u] = timer;
}

bool is_ancestor(int u, int v) {
    return tin[u] <= tin[v] && tout[v] <= tout[u];
}

int lca(int u, int v) {
    if (is_ancestor(u, v)) return u;
    if (is_ancestor(v, u)) return v;
    for (int i = LG - 1; i >= 0; --i)
        if (!is_ancestor(parent[u][i], v))
            u = parent[u][i];
    return parent[u][0];
}

ll get_dist(int u, int v) {
    return dist[u] + dist[v] - 2 * dist[lca(u, v)];
}

vector<pair<int, ll>> vtree[N];
int color[N];  // 0 = none, 1 = group A, 2 = group B
ll dpA[N], dpB[N];

vector<int> build_virtual_tree(vector<int>& nodes) {
    sort(nodes.begin(), nodes.end(), [](int u, int v) {
        return tin[u] < tin[v];
    });

    int sz = nodes.size();
    vector<int> lcas;
    for (int i = 0; i + 1 < sz; ++i) {
        lcas.push_back(lca(nodes[i], nodes[i + 1]));
    }

    for (int x : lcas) nodes.push_back(x);
    sort(nodes.begin(), nodes.end(), [](int u, int v) {
        return tin[u] < tin[v];
    });
    nodes.erase(unique(nodes.begin(), nodes.end()), nodes.end());

    stack<int> st;
    st.push(nodes[0]);
    for (int i = 1; i < (int)nodes.size(); ++i) {
        while (!is_ancestor(st.top(), nodes[i])) st.pop();
        int u = st.top(), v = nodes[i];
        ll w = get_dist(u, v);
        vtree[u].emplace_back(v, w);
        st.push(v);
    }
    return nodes;
}

void dfs_dp(int u) {
    if (color[u] == 1) dpA[u] = 0;
    if (color[u] == 2) dpB[u] = 0;

    for (auto [v, w] : vtree[u]) {
        dfs_dp(v);
        dpA[u] = min(dpA[u], dpA[v] + w);
        dpB[u] = min(dpB[u], dpB[v] + w);
    }
}

void Init(int N_, int A[], int B[], int D[]) {
    n = N_;
    for (int i = 0; i < n - 1; ++i) {
        int u = A[i], v = B[i]; ll w = D[i];
        g[u].emplace_back(v, w);
        g[v].emplace_back(u, w);
    }
    dfs(0, 0);
}

long long Query(int S, int X[], int T, int Y[]) {
    vector<int> nodes;
    for (int i = 0; i < S; ++i) {
        color[X[i]] = 1;
        nodes.push_back(X[i]);
    }
    for (int i = 0; i < T; ++i) {
        color[Y[i]] = 2;
        nodes.push_back(Y[i]);
    }

    vector<int> tree = build_virtual_tree(nodes);

    for (int u : tree) {
        dpA[u] = dpB[u] = INF;
    }

    dfs_dp(tree[0]);

    ll res = INF;
    for (int u : tree) {
        res = min(res, dpA[u] + dpB[u]);
    }

    // cleanup
    for (int u : tree) {
        vtree[u].clear();
        color[u] = 0;
    }

    return res;
}

• Subtask 1: 15 / 15
• Subtask 2: 18 / 18
• Subtask 3: 67 / 67