最小費用流

Share:
  1const inf = 1 << 60
  2
  3type edge struct{ from, to, cap, flow, cost int }
  4type _edge struct{ to, cap, rev, cost int }
  5
  6type minCostFlow struct {
  7	n  int
  8	g  [][]_edge
  9	es [][2]int
 10}
 11
 12func newMinCostFlow(n int) *minCostFlow {
 13	return &minCostFlow{n, make([][]_edge, n), [][2]int{}}
 14}
 15
 16func (mcf *minCostFlow) AddEdge(from, to, cap, cost int) {
 17	mcf.es = append(mcf.es, [2]int{from, len(mcf.g[from])})
 18	mcf.g[from] = append(mcf.g[from], _edge{to, cap, len(mcf.g[to]), cost})
 19	mcf.g[to] = append(mcf.g[to], _edge{from, 0, len(mcf.g[from]) - 1, -cost})
 20}
 21
 22func (mcf *minCostFlow) Edge(i int) edge {
 23	v := mcf.es[i]
 24	e := mcf.g[v[0]][v[1]]
 25	r := mcf.g[e.to][e.rev]
 26	return edge{v[0], e.to, e.cap + r.cap, r.cap, e.cost}
 27}
 28
 29func (mcf *minCostFlow) Edges() []edge {
 30	es := []edge{}
 31	for i := 0; i < len(mcf.es); i++ {
 32		es = append(es, mcf.Edge(i))
 33	}
 34	return es
 35}
 36
 37func (mcf *minCostFlow) Flow(s, t, f int) (int, int) {
 38	n := mcf.n
 39	flow, cost := 0, 0
 40	pv, pe := make([]int, n), make([]int, n)
 41	h, d := make([]int, n), make([]int, n)
 42
 43	min := func(a, b int) int {
 44		if a < b {
 45			return a
 46		}
 47		return b
 48	}
 49
 50	for f > 0 {
 51		for i := 0; i < n; i++ {
 52			d[i] = inf
 53		}
 54		d[s] = 0
 55		pq := &pqForMinCostFlow{}
 56		heap.Push(pq, [2]int{0, s})
 57		for len(*pq) != 0 {
 58			p := heap.Pop(pq).([2]int)
 59			v := p[1]
 60			if d[v] < p[0] {
 61				continue
 62			}
 63			if v == t {
 64				break
 65			}
 66			for i := 0; i < len(mcf.g[v]); i++ {
 67				e := mcf.g[v][i]
 68				if e.cap == 0 {
 69					continue
 70				}
 71				if co := d[v] + e.cost + h[v] - h[e.to]; d[e.to] > co {
 72					d[e.to] = co
 73					pv[e.to], pe[e.to] = v, i
 74					heap.Push(pq, [2]int{d[e.to], e.to})
 75				}
 76			}
 77		}
 78		if d[t] == inf {
 79			return -1, 0
 80		}
 81		for i := 0; i < n; i++ {
 82			h[i] += d[i]
 83		}
 84		d := f
 85		for v := t; v != s; v = pv[v] {
 86			d = min(d, mcf.g[pv[v]][pe[v]].cap)
 87		}
 88		f -= d
 89		flow += d
 90		cost += d * h[t]
 91		for v := t; v != s; v = pv[v] {
 92			e := &mcf.g[pv[v]][pe[v]]
 93			e.cap -= d
 94			mcf.g[v][e.rev].cap += d
 95		}
 96	}
 97
 98	return flow, cost
 99}
100
101type pqForMinCostFlow [][2]int
102
103func (pq pqForMinCostFlow) Len() int            { return len(pq) }
104func (pq pqForMinCostFlow) Swap(i, j int)       { pq[i], pq[j] = pq[j], pq[i] }
105func (pq pqForMinCostFlow) Less(i, j int) bool  { return pq[i][0] < pq[j][0] }
106func (pq *pqForMinCostFlow) Push(x interface{}) { *pq = append(*pq, x.([2]int)) }
107func (pq *pqForMinCostFlow) Pop() interface{} {
108	x := (*pq)[len(*pq)-1]
109	*pq = (*pq)[0 : len(*pq)-1]
110	return x
111}
112