Submission #1163689

#TimeUsernameProblemLanguageResultExecution timeMemory
1163689monkey133Rainy Markets (CCO22_day1problem2)C++20
0 / 25
0 ms328 KiB
#include <iostream>
#include <vector>
#include <queue>
#include <limits>
#include <cassert>
#include <algorithm>
using namespace std;

typedef long long ll;
const ll INF = 1e18;

struct Edge {
    int to, rev;
    ll cap, cost;
    ll orig; // record original capacity (for assignment extraction)
};

struct MinCostFlow {
    int n;
    vector<vector<Edge>> graph;
    vector<ll> dist, potential;
    vector<int> parent, parentEdge;
    
    MinCostFlow(int n): n(n), graph(n), dist(n), potential(n), parent(n), parentEdge(n) {}
    
    // Add a directed edge from s to t with given capacity and cost.
    void addEdge(int s, int t, ll cap, ll cost) {
        Edge a = { t, (int)graph[t].size(), cap, cost, cap };
        Edge b = { s, (int)graph[s].size(), 0, -cost, 0 };
        graph[s].push_back(a);
        graph[t].push_back(b);
    }
    
    // Sends at most f units from s to t.
    // Returns a pair {flow sent, total cost}.
    pair<ll, ll> minCostFlow(int s, int t, ll f) {
        ll flow = 0, cost = 0;
        // initialize potentials to 0
        fill(potential.begin(), potential.end(), 0);
        while (flow < f) {
            // Use Dijkstra to find shortest augmenting path (w.r.t. adjusted costs)
            fill(dist.begin(), dist.end(), INF);
            dist[s] = 0;
            priority_queue<pair<ll,int>, vector<pair<ll,int>>, greater<pair<ll,int>>> pq;
            pq.push({0, s});
            while(!pq.empty()){
                auto [d, v] = pq.top();
                pq.pop();
                if(d != dist[v]) continue;
                for (int i = 0; i < graph[v].size(); i++){
                    Edge &e = graph[v][i];
                    if(e.cap > 0 && dist[e.to] > dist[v] + e.cost + potential[v] - potential[e.to]){
                        dist[e.to] = dist[v] + e.cost + potential[v] - potential[e.to];
                        parent[e.to] = v;
                        parentEdge[e.to] = i;
                        pq.push({dist[e.to], e.to});
                    }
                }
            }
            if(dist[t] == INF) break; // no more augmenting path; flow not possible
            for (int i = 0; i < n; i++){
                if(dist[i] < INF)
                    potential[i] += dist[i];
            }
            ll addFlow = f - flow;
            for (int v = t; v != s; v = parent[v]){
                addFlow = min(addFlow, graph[parent[v]][parentEdge[v]].cap);
            }
            flow += addFlow;
            cost += addFlow * potential[t];
            for (int v = t; v != s; v = parent[v]){
                Edge &e = graph[parent[v]][parentEdge[v]];
                e.cap -= addFlow;
                graph[v][e.rev].cap += addFlow;
            }
        }
        return {flow, cost};
    }
};

//
// Main function.
// Input format:
//   First line: integer N (number of bus shelters)
//   Second line: N integers B[1..N] (capacity of each bus shelter)
//   Third line: N-1 integers P[1..N-1] (people at market i)
//   Fourth line: N-1 integers U[1..N-1] (umbrellas available at market i)
// 
// Output format:
//   If a valid assignment exists, output N+1 lines:
//      YES
//      <minimum umbrella cost>
//      Then for each market i from 1 to N-1, a line with three space–separated integers:
//          <# going to bus shelter i> <# buying umbrella> <# going to bus shelter i+1>
//   If not, output a single line: NO
//
int main(){
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    
    int N;
    cin >> N;
    // bus shelters are numbered 1..N; read capacities B[1..N]
    vector<ll> B(N+1);
    for (int i = 1; i <= N; i++){
        cin >> B[i];
    }
    
    int M = N - 1; // markets: 1 .. N-1
    vector<ll> P(M+1), U(M+1);
    for (int i = 1; i <= M; i++){
        cin >> P[i];
    }
    for (int i = 1; i <= M; i++){
        cin >> U[i];
    }
    
    // We build a network with the following nodes:
    // S (source) = 0.
    // Market nodes: 1 .. M.
    // Bus shelter nodes: M+1 .. M+N.
    // T (sink) = M+N+1.
    int S = 0;
    int marketStart = 1;
    int busStart = marketStart + M; // bus shelter nodes start here; bus shelter i is at node (busStart + i - 1)
    int T = busStart + N;
    int totalNodes = T + 1;
    
    MinCostFlow mcf(totalNodes);
    
    // Total number of people in all markets.
    ll totalPeople = 0;
    for (int i = 1; i <= M; i++){
        totalPeople += P[i];
    }
    
    // 1. From source S to each market i with capacity P[i] and cost 0.
    for (int i = 1; i <= M; i++){
        mcf.addEdge(S, i, P[i], 0);
    }
    
    // 2. For each market i, add three edges:
    //    a) Market i -> bus shelter i ("left" option). (Node: busStart + (i-1))
    //    b) Market i -> bus shelter i+1 ("right" option). (Node: busStart + i)
    //    c) Market i -> sink T (buy umbrella) with capacity U[i] and cost 1.
    for (int i = 1; i <= M; i++){
        mcf.addEdge(i, busStart + (i - 1), P[i], 0); // left option
        mcf.addEdge(i, busStart + i, P[i], 0);         // right option
        mcf.addEdge(i, T, U[i], 1);                      // umbrella option
    }
    
    // 3. For each bus shelter i (1..N), add edge from bus shelter node to sink T with capacity B[i] and cost 0.
    for (int i = 1; i <= N; i++){
        mcf.addEdge(busStart + (i - 1), T, B[i], 0);
    }
    
    auto res = mcf.minCostFlow(S, T, totalPeople);
    
    if(res.first < totalPeople) {
        cout << "NO\n";
        return 0;
    }
    
    // If we reached here, a valid assignment exists.
    cout << "YES\n" << res.second << "\n";
    
    // For each market i (nodes 1..M), we extract the flow that left along its three edges.
    // Recall that we added the edges in the following order for each market i:
    //   index 0: Market i -> bus shelter i ("left")
    //   index 1: Market i -> bus shelter i+1 ("right")
    //   index 2: Market i -> sink T (umbrella)
    // The flow used on an edge = (original capacity) - (remaining capacity).
    for (int i = 1; i <= M; i++){
        assert(mcf.graph[i].size() >= 3);
        ll left_flow = mcf.graph[i][0].orig - mcf.graph[i][0].cap;
        ll right_flow = mcf.graph[i][1].orig - mcf.graph[i][1].cap;
        ll umb_flow = mcf.graph[i][2].orig - mcf.graph[i][2].cap;
        // Output: <people going to bus shelter i> <people buying umbrella> <people going to bus shelter i+1>
        cout << left_flow << " " << umb_flow << " " << right_flow << "\n";
    }
    
    return 0;
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...