This documentation is automatically generated by online-judge-tools/verification-helper
src/min_cost_flow.cr
- View this file on GitHub
- Last update: 2023-11-28 17:27:51+09:00
- Link: View error logs on GitHub Actions
Depends on
Required by
spec/min_cost_flow_spec.cr
Verified with
Code
# ac-library.cr by hakatashi https://github.com/google/ac-library.cr
#
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
require "./graph.cr"
module AtCoder
# Implements [atcoder::mcf_graph](https://atcoder.github.io/ac-library/master/document_en/mincostflow.html).
#
# ```
# flow = AtCoder::MinCostFlow.new(5)
# flow.add_edge(0, 1, 30, 3)
# flow.add_edge(0, 2, 60, 9)
# flow.add_edge(1, 2, 40, 5)
# flow.add_edge(1, 3, 50, 7)
# flow.add_edge(2, 3, 20, 8)
# flow.add_edge(2, 4, 50, 6)
# flow.add_edge(3, 4, 60, 7)
# flow.flow(0, 4, 70) # => {70, 1080}
# ```
class MinCostFlow
private record EdgeInfo, capacity : Int64, cost : Int64 | Nil do
def self.zero
new(Int64::MAX, 0_i64)
end
def +(edge : EdgeInfo)
from_cost = @cost
to_cost = edge.cost
if from_cost.nil? || to_cost.nil?
return self.class.new(0_i64, nil)
end
self.class.new(min(@capacity, edge.capacity), from_cost + to_cost)
end
def >(edge : EdgeInfo)
from_dist = dist
to_dist = edge.dist
return true if from_dist.nil?
return false if to_dist.nil?
from_dist > to_dist
end
def dist
return nil if @capacity == 0
@cost
end
@[AlwaysInline]
private def min(a, b)
a > b ? b : a
end
end
# Implements atcoder::mcf_graph g(n).
def initialize(@size : Int64)
@graph = AtCoder::DirectedGraph(Nil, EdgeInfo).new(@size)
@dists = Array(Int64 | Nil).new(@size, 0_i64)
@edges = [] of {Int64, Int64}
end
# Implements atcoder::mcf_graph.add_edge(from, to, capacity, cost).
def add_edge(from, to, capacity, cost)
@graph.add_edge(from, to, EdgeInfo.new(capacity.to_i64, cost.to_i64))
@graph.add_edge(to, from, EdgeInfo.new(0_i64, -cost.to_i64))
@edges << {from.to_i64, to.to_i64}
end
private def increment_edge_cost(from, to, cost_diff)
edge = @graph.get_edge(from, to)
@graph.update_edge(from, to, edge + EdgeInfo.new(edge.capacity, cost_diff))
end
private def increment_edge_capacity(from, to, capacity_diff)
edge = @graph.get_edge(from, to)
@graph.update_edge(from, to, EdgeInfo.new(edge.capacity + capacity_diff, edge.cost))
end
# Implements atcoder::mcf_graph.slope(start, target, flow_limit).
# ameba:disable Metrics/CyclomaticComplexity
def slope(start, target, flow_limit : Int | Nil = nil)
raise ArgumentError.new("start and target cannot be the same") if start == target
flow_points = [] of {Int64, Int64}
current_cost = 0_i64
flowed_capacity = 0_i64
min_cost = 0_i64
until flowed_capacity == flow_limit
nodes = @graph.dijkstra(start)
target_dist = nodes[target][:dist]
if target_dist.nil? || target_dist.capacity == 0
break
end
capacity = target_dist.capacity
# Update edge capacities on the path from start to target
last_node = target
until last_node == start
prev_node = nodes[last_node][:prev].not_nil!
increment_edge_capacity(prev_node, last_node, -capacity)
increment_edge_capacity(last_node, prev_node, capacity)
last_node = prev_node
end
# Update edge costs
@edges.each do |from, to|
from_dist = nodes[from][:dist]
to_dist = nodes[to][:dist]
next if from_dist.nil? || to_dist.nil?
next if from_dist.cost.nil? || to_dist.cost.nil?
dist = to_dist.cost.not_nil! - from_dist.cost.not_nil!
increment_edge_cost(from, to, -dist)
increment_edge_cost(to, from, dist)
end
# Update distants
nodes.each_with_index do |node, i|
dist = node[:dist]
if dist.nil? || @dists[i].nil? || dist.not_nil!.cost.nil?
@dists[i] = nil
else
@dists[i] = @dists[i].not_nil! + dist.not_nil!.cost.not_nil!
end
end
new_cost = @dists[target].not_nil!
if flow_limit.nil?
new_capacity = capacity
else
new_capacity = min(capacity, flow_limit - flowed_capacity)
end
if new_cost != current_cost
if current_cost == 0 && flowed_capacity != 0
flow_points << {0_i64, 0_i64}
end
flow_points << {flowed_capacity, min_cost}
end
min_cost += new_cost * new_capacity
flowed_capacity += new_capacity
current_cost = new_cost
end
flow_points << {flowed_capacity, min_cost}
flow_points
end
# Implements atcoder::mcf_graph.flow(start, target, flow_limit).
def flow(start, target, flow_limit : Int | Nil = nil)
flow_points = slope(start, target, flow_limit)
flow_points.last
end
@[AlwaysInline]
private def min(a, b)
a > b ? b : a
end
delegate :size, to: @graph
end
end# ac-library.cr by hakatashi https://github.com/google/ac-library.cr
#
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# require "./graph.cr"
# ac-library.cr by hakatashi https://github.com/google/ac-library.cr
#
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# require "./priority_queue.cr"
# ac-library.cr by hakatashi https://github.com/google/ac-library.cr
#
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
module AtCoder
# Implements standard priority queue like [std::priority_queue](https://en.cppreference.com/w/cpp/container/priority_queue).
#
# ```
# q = AtCoder::PriorityQueue(Int64).new
# q << 1_i64
# q << 3_i64
# q << 2_i64
# q.pop # => 3
# q.pop # => 2
# q.pop # => 1
# ```
class PriorityQueue(T)
include Enumerable(T)
getter heap : Array(T)
# Create a new queue in ascending order of priority.
def self.max
self.new { |a, b| a <= b }
end
# Create a new queue in ascending order of priority with the elements in *enumerable*.
def self.max(enumerable : Enumerable(T))
self.new(enumerable) { |a, b| a <= b }
end
# Create a new queue in descending order of priority.
def self.min
self.new { |a, b| a >= b }
end
# Create a new queue in descending order of priority with the elements in *enumerable*.
def self.min(enumerable : Enumerable(T))
self.new(enumerable) { |a, b| a >= b }
end
def initialize
initialize { |a, b| a <= b }
end
# Initializes queue with the elements in *enumerable*.
def self.new(enumerable : Enumerable(T))
self.new(enumerable) { |a, b| a <= b }
end
# Initializes queue with the custom comperator.
#
# If the second argument `b` should be popped earlier than
# the first argument `a`, return `true`. Else, return `false`.
#
# ```
# q = AtCoder::PriorityQueue(Int64).new { |a, b| a >= b }
# q << 1_i64
# q << 3_i64
# q << 2_i64
# q.pop # => 1
# q.pop # => 2
# q.pop # => 3
# ```
def initialize(&block : T, T -> Bool)
@heap = Array(T).new
@compare_proc = block
end
# Initializes queue with the elements in *enumerable* and the custom comperator.
#
# If the second argument `b` should be popped earlier than
# the first argument `a`, return `true`. Else, return `false`.
#
# ```
# q = AtCoder::PriorityQueue.new([1, 3, 2]) { |a, b| a >= b }
# q.pop # => 1
# q.pop # => 2
# q.pop # => 3
# ```
def initialize(enumerable : Enumerable(T), &block : T, T -> Bool)
@heap = enumerable.to_a
@compare_proc = block
len = @heap.size
(len // 2 - 1).downto(0) do |parent|
v = @heap[parent]
child = parent * 2 + 1
while child < len
if child + 1 < len && @compare_proc.call(@heap[child], @heap[child + 1])
child += 1
end
unless @compare_proc.call(v, @heap[child])
break
end
@heap[parent] = @heap[child]
parent = child
child = parent * 2 + 1
end
@heap[parent] = v
end
end
# Pushes value into the queue.
# This method returns self, so several calls can be chained.
def push(v : T) : self
@heap << v
index = @heap.size - 1
while index != 0
parent = (index - 1) // 2
if @compare_proc.call(@heap[index], @heap[parent])
break
end
@heap[parent], @heap[index] = @heap[index], @heap[parent]
index = parent
end
self
end
# Alias of `push`
def <<(v : T) : self
push(v)
end
# Pops value from the queue.
def pop
if @heap.size == 0
return nil
end
if @heap.size == 1
return @heap.pop
end
ret = @heap.first
@heap[0] = @heap.pop
index = 0
while index * 2 + 1 < @heap.size
child = if index * 2 + 2 < @heap.size && !@compare_proc.call(@heap[index * 2 + 2], @heap[index * 2 + 1])
index * 2 + 2
else
index * 2 + 1
end
if @compare_proc.call(@heap[child], @heap[index])
break
end
@heap[child], @heap[index] = @heap[index], @heap[child]
index = child
end
ret
end
# Pops value from the queue.
# Raises `Enumerable::EmptyError` if queue is of 0 size.
def pop!
pop || raise Enumerable::EmptyError.new
end
# Yields each item in the queue in comparator's order.
def each(&)
@heap.sort { |a, b| @compare_proc.call(a, b) ? 1 : -1 }.each do |e|
yield e
end
end
# Returns, but does not remove, the head of the queue.
def first(&)
@heap.first { yield }
end
# Returns `true` if the queue is empty.
delegate :empty?, to: @heap
# Returns size of the queue.
delegate :size, to: @heap
end
end
module AtCoder
class Graph(NodeInfo, EdgeInfo)
@size_bits : Int32
getter visited : Set(Int64)
def initialize(@nodes : Array(NodeInfo))
@size = @nodes.size.to_i64
@edges = [] of EdgeInfo
@adjacencies = Array(Hash(Int64, Int64)).new(@size) { Hash(Int64, Int64).new }
@visited = Set(Int64).new
@size_bits = @size.to_s(2).size
end
def initialize(@size : Int64, initial_node : NodeInfo = nil)
@nodes = Array(NodeInfo).new(@size, initial_node)
@edges = [] of EdgeInfo
@adjacencies = Array(Hash(Int64, Int64)).new(@size) { Hash(Int64, Int64).new }
@visited = Set(Int64).new
@size_bits = @size.to_s(2).size
end
# Performs Dijkstra's Algorithm to calculate the distance of each node from `start_node`.
# To use this method, `EdgeInfo` must implement `.zero` and `#+(EdgeInfo)` and `#>(EdgeInfo)`.
def dijkstra(start_node)
dist = Array(EdgeInfo | Nil).new(@size, nil)
dist[start_node] = EdgeInfo.zero
prev_nodes = Array(Int64 | Nil).new(@size, nil)
queue = AtCoder::PriorityQueue({EdgeInfo, Int64}).new { |(ad, av), (bd, bv)| ad > bd }
queue << {EdgeInfo.zero, start_node.to_i64}
until queue.empty?
d, v = queue.pop.not_nil!
current_dist = dist[v].not_nil!
next if d > current_dist
@adjacencies[v].each do |(to, edge_id)|
cost = @edges[edge_id]
target_dist = dist[to]
if target_dist.nil? || target_dist > current_dist + cost
dist[to] = current_dist + cost
prev_nodes[to] = v
queue << {current_dist + cost, to}
end
end
end
Array({dist: EdgeInfo | Nil, prev: Int64 | Nil}).new(@size) do |i|
{dist: dist[i], prev: prev_nodes[i]}
end
end
def dfs(node : Int64, initial_value : T, &block : Int64, T, NamedTuple(
node: Int64,
node_info: NodeInfo | Nil,
edge: Int64 | Nil,
edge_info: EdgeInfo | Nil,
parent: Int64,
), (T ->) ->) forall T
@visited = Set(Int64).new
info = {
node: node,
node_info: @nodes[node].as(NodeInfo | Nil),
edge: nil.as(Int64 | Nil),
edge_info: nil.as(EdgeInfo | Nil),
parent: -1_i64,
}
block.call(node, initial_value, info, ->(new_value : T) {
dfs(node, -1_i64, new_value, &block)
})
end
private def dfs(node : Int64, parent : Int64, value : T, &block : Int64, T, NamedTuple(
node: Int64,
node_info: NodeInfo | Nil,
edge: Int64 | Nil,
edge_info: EdgeInfo | Nil,
parent: Int64,
), (T ->) ->) forall T
@visited << node
@adjacencies[node].each do |(child, edge)|
next if @visited.includes?(child)
info = {
node: child,
node_info: @nodes[child].as(NodeInfo | Nil),
edge: edge.as(Int64 | Nil),
edge_info: @edges[edge].as(EdgeInfo | Nil),
parent: node,
}
block.call(child, value, info, ->(new_value : T) {
dfs(child, node, new_value, &block)
})
end
end
end
class DirectedGraph(NodeInfo, EdgeInfo) < Graph(NodeInfo, EdgeInfo)
def add_edge(from, to, edge : EdgeInfo = 1_i64)
@edges << edge
edge_id = @edges.size.to_i64 - 1
@adjacencies[from][to.to_i64] = edge_id
edge_id
end
def get_edge(from, to)
edge_id = @adjacencies[from][to.to_i64]
@edges[edge_id]
end
def update_edge(from, to, edge : EdgeInfo = 1_i64)
edge_id = @adjacencies[from][to.to_i64]
@edges[edge_id] = edge
edge_id
end
end
class UndirectedGraph(NodeInfo, EdgeInfo) < Graph(NodeInfo, EdgeInfo)
def add_edge(a, b, edge : EdgeInfo = 1_i64)
@edges << edge
edge_id = @edges.size.to_i64 - 1
@adjacencies[a][b.to_i64] = edge_id
@adjacencies[b][a.to_i64] = edge_id
edge_id
end
def get_edge(a, b)
edge_id = @adjacencies[a][b.to_i64]
@edges[edge_id]
end
def update_edge(a, b, edge : EdgeInfo = 1_i64)
edge_id = @adjacencies[a][b.to_i64]
@edges[edge_id] = edge
edge_id
end
end
class Tree(NodeInfo, EdgeInfo) < UndirectedGraph(NodeInfo, EdgeInfo)
@lca_root : Int64 | Nil
@lca_ancestors : Array(Array(Int64)) | Nil
@lca_depths : Array(Int64) | Nil
def diameter
@farthest_node = -1_i64
@farthest_depth = 0_i64
dfs(0_i64, 0_i64) do |node, depth, info, callback|
weight = info[:edge_info]
depth += weight.nil? ? 0 : weight
if @farthest_depth.not_nil! < depth
@farthest_node = node
@farthest_depth = depth
end
callback.call(depth)
end
start_node = @farthest_node.not_nil!
@farthest_node = -1_i64
@farthest_depth = 0_i64
@parents = Array(Int64).new(@size, -1_i64)
dfs(start_node, 0_i64) do |node, depth, info, callback|
weight = info[:edge_info]
depth += weight.nil? ? 0 : weight
@parents.not_nil![node] = info[:parent]
if @farthest_depth.not_nil! < depth
@farthest_node = node
@farthest_depth = depth
end
callback.call(depth)
end
{@farthest_depth.not_nil!, start_node, @farthest_node.not_nil!, @parents.not_nil!}
end
private def lca_precompute(root)
lca_ancestors = Array(Array(Int64)).new(@size) { Array(Int64).new(@size_bits, -1_i64) }
lca_depths = Array(Int64).new(@size, -1_i64)
dfs(root, 0_i64) do |node, depth, info, callback|
lca_ancestors[node][0] = info[:parent]
lca_depths[node] = depth
callback.call(depth + 1)
end
1.upto(@size_bits - 1) do |i|
@size.times do |node|
if lca_ancestors[node][i - 1] == -1
lca_ancestors[node][i] = -1_i64
else
lca_ancestors[node][i] = lca_ancestors[lca_ancestors[node][i - 1]][i - 1]
end
end
end
@lca_root = root
@lca_ancestors = lca_ancestors
@lca_depths = lca_depths
end
private def lca_nth_prev(node, dist)
lca_ancestors = @lca_ancestors.not_nil!
i = 0_i64
until dist == 0
if dist.odd?
node = lca_ancestors[node][i]
if node == -1
raise Exception.new("#{dist}th previous node of #{node} does not exist")
end
end
dist //= 2
i += 1
end
node
end
# ameba:disable Metrics/CyclomaticComplexity
def lca(a, b)
if @lca_root.nil?
lca_precompute(0_i64)
end
lca_ancestors = @lca_ancestors.not_nil!
lca_depths = @lca_depths.not_nil!
depth_a = lca_depths[a]
depth_b = lca_depths[b]
if depth_a < depth_b
depth_a, depth_b, a, b = depth_b, depth_a, b, a
end
if depth_a != depth_b
a = lca_nth_prev(a, depth_a - depth_b)
end
if a == b
return a
end
(@size_bits - 1).downto(0) do |i|
if lca_ancestors[a][i] == -1 || lca_ancestors[b][i] == -1
next
end
if lca_ancestors[a][i] != lca_ancestors[b][i]
a = lca_ancestors[a][i]
b = lca_ancestors[b][i]
end
end
a = lca_ancestors[a][0]
b = lca_ancestors[b][0]
if a != b || a == -1 || b == -1
raise Exception.new("Assertion error")
end
a
end
def dist(a, b)
lca_node = lca(a, b)
lca_depths = @lca_depths.not_nil!
(lca_depths[a] - lca_depths[lca_node]) + (lca_depths[b] - lca_depths[lca_node])
end
# Returns `true` if node `c` is on the path from `a` to `b`.
def on_path?(a, b, c)
dist(a, c) + dist(c, b) == dist(a, b)
end
end
end
module AtCoder
# Implements [atcoder::mcf_graph](https://atcoder.github.io/ac-library/master/document_en/mincostflow.html).
#
# ```
# flow = AtCoder::MinCostFlow.new(5)
# flow.add_edge(0, 1, 30, 3)
# flow.add_edge(0, 2, 60, 9)
# flow.add_edge(1, 2, 40, 5)
# flow.add_edge(1, 3, 50, 7)
# flow.add_edge(2, 3, 20, 8)
# flow.add_edge(2, 4, 50, 6)
# flow.add_edge(3, 4, 60, 7)
# flow.flow(0, 4, 70) # => {70, 1080}
# ```
class MinCostFlow
private record EdgeInfo, capacity : Int64, cost : Int64 | Nil do
def self.zero
new(Int64::MAX, 0_i64)
end
def +(edge : EdgeInfo)
from_cost = @cost
to_cost = edge.cost
if from_cost.nil? || to_cost.nil?
return self.class.new(0_i64, nil)
end
self.class.new(min(@capacity, edge.capacity), from_cost + to_cost)
end
def >(edge : EdgeInfo)
from_dist = dist
to_dist = edge.dist
return true if from_dist.nil?
return false if to_dist.nil?
from_dist > to_dist
end
def dist
return nil if @capacity == 0
@cost
end
@[AlwaysInline]
private def min(a, b)
a > b ? b : a
end
end
# Implements atcoder::mcf_graph g(n).
def initialize(@size : Int64)
@graph = AtCoder::DirectedGraph(Nil, EdgeInfo).new(@size)
@dists = Array(Int64 | Nil).new(@size, 0_i64)
@edges = [] of {Int64, Int64}
end
# Implements atcoder::mcf_graph.add_edge(from, to, capacity, cost).
def add_edge(from, to, capacity, cost)
@graph.add_edge(from, to, EdgeInfo.new(capacity.to_i64, cost.to_i64))
@graph.add_edge(to, from, EdgeInfo.new(0_i64, -cost.to_i64))
@edges << {from.to_i64, to.to_i64}
end
private def increment_edge_cost(from, to, cost_diff)
edge = @graph.get_edge(from, to)
@graph.update_edge(from, to, edge + EdgeInfo.new(edge.capacity, cost_diff))
end
private def increment_edge_capacity(from, to, capacity_diff)
edge = @graph.get_edge(from, to)
@graph.update_edge(from, to, EdgeInfo.new(edge.capacity + capacity_diff, edge.cost))
end
# Implements atcoder::mcf_graph.slope(start, target, flow_limit).
# ameba:disable Metrics/CyclomaticComplexity
def slope(start, target, flow_limit : Int | Nil = nil)
raise ArgumentError.new("start and target cannot be the same") if start == target
flow_points = [] of {Int64, Int64}
current_cost = 0_i64
flowed_capacity = 0_i64
min_cost = 0_i64
until flowed_capacity == flow_limit
nodes = @graph.dijkstra(start)
target_dist = nodes[target][:dist]
if target_dist.nil? || target_dist.capacity == 0
break
end
capacity = target_dist.capacity
# Update edge capacities on the path from start to target
last_node = target
until last_node == start
prev_node = nodes[last_node][:prev].not_nil!
increment_edge_capacity(prev_node, last_node, -capacity)
increment_edge_capacity(last_node, prev_node, capacity)
last_node = prev_node
end
# Update edge costs
@edges.each do |from, to|
from_dist = nodes[from][:dist]
to_dist = nodes[to][:dist]
next if from_dist.nil? || to_dist.nil?
next if from_dist.cost.nil? || to_dist.cost.nil?
dist = to_dist.cost.not_nil! - from_dist.cost.not_nil!
increment_edge_cost(from, to, -dist)
increment_edge_cost(to, from, dist)
end
# Update distants
nodes.each_with_index do |node, i|
dist = node[:dist]
if dist.nil? || @dists[i].nil? || dist.not_nil!.cost.nil?
@dists[i] = nil
else
@dists[i] = @dists[i].not_nil! + dist.not_nil!.cost.not_nil!
end
end
new_cost = @dists[target].not_nil!
if flow_limit.nil?
new_capacity = capacity
else
new_capacity = min(capacity, flow_limit - flowed_capacity)
end
if new_cost != current_cost
if current_cost == 0 && flowed_capacity != 0
flow_points << {0_i64, 0_i64}
end
flow_points << {flowed_capacity, min_cost}
end
min_cost += new_cost * new_capacity
flowed_capacity += new_capacity
current_cost = new_cost
end
flow_points << {flowed_capacity, min_cost}
flow_points
end
# Implements atcoder::mcf_graph.flow(start, target, flow_limit).
def flow(start, target, flow_limit : Int | Nil = nil)
flow_points = slope(start, target, flow_limit)
flow_points.last
end
@[AlwaysInline]
private def min(a, b)
a > b ? b : a
end
delegate :size, to: @graph
end
end