제출 #1150960

#	제출 시각	아이디	문제	언어	결과	실행 시간	메모리
1150960		arshjeet	Toll (APIO13_toll)	C++20	16 / 100	1 ms	2628 KiB

toll

#include <bits/stdc++.h>
using namespace std;
#define int long long
#define pb push_back
#define fast() ios_base::sync_with_stdio(false); cin.tie(NULL)

const int N = 1e5 + 1;
const int INF = 1e6 + 1;
set<vector<int>> edges; // wt, v1, v2
vector<pair<int, int>> mst[N];
int sz[N], parent[N], level[N], bfsvis[N], ppl[N];

// DFS to count participants passing through each node
int dfs(int v, int p) {
    int ans = ppl[v];
    for (auto node : mst[v]) {
        int child = node.first;
        if (child == p) continue;
        ans += dfs(child, v);
    }
    return ans;
}

// BFS to compute levels of nodes relative to town 1
void bfs(int src) {
    queue<int> q;
    level[src] = 0;
    bfsvis[src] = 1;
    q.push(src);

    while (!q.empty()) {
        int v = q.front();
        q.pop();
        for (auto node : mst[v]) {
            int child = node.first;
            if (!bfsvis[child]) {
                bfsvis[child] = 1;
                level[child] = level[v] + 1;
                q.push(child);
            }
        }
    }
}

// Union-Find Data Structure
void make(int a) {
    parent[a] = a;
    sz[a] = 1;
}

int find(int v) {
    if (v == parent[v]) return v;
    return parent[v] = find(parent[v]);
}

void Union(int a, int b) {
    a = find(a);
    b = find(b);
    if (a != b) {
        if (sz[a] < sz[b]) swap(a, b);
        parent[b] = a;
        sz[a] += sz[b];
    }
}

// Kruskal's Algorithm to construct MST
void kruskal(int n) {
    for (int i = 1; i <= n; ++i) make(i);
    for (int i = 1; i <= n; ++i) mst[i].clear();

    for (auto edge : edges) {
        int wt = edge[0], v1 = edge[1], v2 = edge[2];
        if (find(v1) == find(v2)) continue;
        Union(v1, v2);
        mst[v1].pb({v2, wt});
        mst[v2].pb({v1, wt});
    }
}

signed main() {
    fast();
    int n, m, k;
    cin >> n >> m >> k;
    map<pair<int, int>, int> dist;

    // Edge case handling
    if (n == 1 || k == 0 || m == 0) {
        cout << 0;
        return 0;
    }

    // Input old roads
    for (int i = 0; i < m; ++i) {
        int a, b, c;
        cin >> a >> b >> c;
        edges.insert({c, a, b});
    }

    // Input new roads
    vector<pair<int, int>> newroads;
    for (int i = 0; i < k; ++i) {
        int x, y;
        cin >> x >> y;
        newroads.pb({x, y});
        edges.insert({INF, x, y});
        dist[{x, y}] = INF;
        dist[{y, x}] = INF;
    }

    // Input participants
    for (int i = 1; i <= n; ++i) cin >> ppl[i];

    // Precompute the MST with old roads
    kruskal(n);

    // Binary search for each new road to find maximum toll fee
    for (int i = 0; i < k; ++i) {
        int v1 = newroads[i].first, v2 = newroads[i].second;

        int lo = 1, hi = INF - 1;
        while (lo < hi) {
            int mid = (lo + hi + 1) / 2;

            // Temporarily add the new road with toll fee = mid
            edges.insert({mid, v1, v2});

            // Reconstruct MST
            kruskal(n);

            // Check if the new road is included in the MST
            bool included = false;
            for (auto node : mst[v1]) {
                if (node.first == v2) {
                    included = true;
                    break;
                }
            }

            if (included) lo = mid;
            else hi = mid - 1;

            // Remove the temporary edge
            edges.erase({mid, v1, v2});
        }

        // Add the new road with the final toll fee
        edges.insert({lo, v1, v2});
        dist[{v1, v2}] = lo;
        dist[{v2, v1}] = lo;
    }

    // Construct the final MST with optimized toll fees
    kruskal(n);

    // Compute levels using BFS
    bfs(1);

    // Calculate revenue from new roads
    int ans = 0;
    for (auto it = dist.begin(); it != dist.end(); ++it) {
        auto pr = it->first;
        int v1 = pr.first, v2 = pr.second;

        // Ensure v2 is farther from town 1 than v1
        if (level[v1] > level[v2]) continue;

        // Count participants passing through the new road
        int z = dfs(v2, v1);
        ans += z * dist[{v1, v2}];
    }

    cout << ans;
}

• Subtask 1: 16 / 16
• Subtask 2: 0 / 18
• Subtask 3: 0 / 22
• Subtask 4: 0 / 22
• Subtask 5: 0 / 22