AtCoder Beginner Contest 004 D: マーブル (最小費用流)

問題

abc004.contest.atcoder.jp

解法

www.slideshare.net

スタートと赤・緑・青のノードを結び、赤・緑・青と座標-500~500のノードを結び、これらの座標とゴールを結んだグラフを作る。最小費用流のクラスは蟻本の写経。

コード

import java.util.ArrayList;
import java.util.Arrays;
import java.util.PriorityQueue;
import java.util.Scanner;

public class Main {

	public static void main(String[] args) {
		Scanner sc = new Scanner(System.in);
		int R = sc.nextInt();
		int G = sc.nextInt();
		int B = sc.nextInt();
		sc.close();

		PrimalDual primalDual = new PrimalDual(1001 + 1 + 1 + 3);
		primalDual.addEdge(0, 1, R, 0);
		primalDual.addEdge(0, 2, G, 0);
		primalDual.addEdge(0, 3, B, 0);
		for (int i = -500; i < 500; i++) {
			primalDual.addEdge(1, i + 504, 1, Math.abs(i + 100));// 赤(-100)
			primalDual.addEdge(2, i + 504, 1, Math.abs(i));// 緑(0)
			primalDual.addEdge(3, i + 504, 1, Math.abs(i - 100));// 青(100)
			primalDual.addEdge(i + 504, 1005, 1, 0);
		}

		System.out.println(primalDual.minCostFlow(0, 1005, R + G + B));

	}

	static class PrimalDual {
		final int INF = 1 << 29;
		int N;// 頂点数
		ArrayList<Edge>[] graph;

		@SuppressWarnings("unchecked")
		public PrimalDual(int N) {
			this.N = N;
			this.graph = new ArrayList[N];
			for (int i = 0; i < N; i++) {
				graph[i] = new ArrayList<Edge>();
			}
		}

		public void addEdge(int from, int to, int cap, int cost) {
			graph[from].add(new Edge(to, cap, cost, graph[to].size()));
			graph[to].add(new Edge(from, 0, -cost, graph[from].size() - 1));
		}

		public int minCostFlow(int start, int goal, int flow) {
			int[] prevNode = new int[N];
			int[] prevEdge = new int[N];
			int[] potential = new int[N];// ポテンシャル(既にかかったコスト)
			int totalCost = 0;
			while (flow > 0) {
				int[] dist = new int[N];
				Arrays.fill(dist, INF);
				dist[start] = 0;

				PriorityQueue<Node> priorityQueue = new PriorityQueue<Node>();
				priorityQueue.offer(new Node(0, start));
				while (!priorityQueue.isEmpty()) {
					// キューから1番距離の近いノードを取り出す
					Node node = priorityQueue.poll();
					int v = node.id;
					if (dist[v] < node.dist) {
						// 暫定の最短距離よりも遠かったらスルー
						continue;
					}

					for (int i = 0; i < graph[v].size(); i++) {
						Edge e = graph[v].get(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];
							priorityQueue.add(new Node(dist[e.to], e.to));

							// 直前の経路を記憶しておく
							prevNode[e.to] = v;
							prevEdge[e.to] = i;
						}
					}
				}

				// そもそもゴールまで辿りつけない場合はどうしようもない
				if (dist[goal] == INF) {
					return -1;
				}

				// 今回かかったコストをメモしておく
				for (int v = 0; v < N; v++) {
					potential[v] += dist[v];
				}

				// startからgoalまで流せるだけ流す
				int minFlow = flow;
				for (int now = goal; now != start; now = prevNode[now]) {
					minFlow = Math.min(minFlow, graph[prevNode[now]].get(prevEdge[now]).cap);
				}
				flow -= minFlow;
				totalCost += minFlow * potential[goal];

				for (int now = goal; now != start; now = prevNode[now]) {
					Edge edge = graph[prevNode[now]].get(prevEdge[now]);
					edge.cap -= minFlow;
					graph[now].get(edge.rev).cap += minFlow;
				}
			}
			return totalCost;
		}

		class Node implements Comparable<Node> {
			int dist;
			int id;

			public Node(int dist, int i) {
				this.dist = dist;
				this.id = i;
			}

			public int compareTo(Node o) {
				return this.dist - o.dist;
			}
		}

		class Edge {
			int to, cap, cost, rev;

			public Edge(int to, int cap, int cost, int rev) {
				this.to = to;
				this.cap = cap;
				this.cost = cost;
				this.rev = rev;
			}
		}
	}

}