This documentation is automatically generated by online-judge-tools/verification-helper
src/graph/re_rooting.cr
- View this file on GitHub
- Last update: 2022-04-16 16:23:13+09:00
Depends on
Verified with
Code
require "../graph"
# Example of `T`:
# ```
# struct DP
# getter val : Int64, cnt : Int32
#
# def initialize
# @val, @cnt = 0i64, 0
# end
#
# def initialize(@val, @cnt)
# end
#
# def +(other : self) : self
# DP.new(val + other.val, cnt + other.cnt)
# end
#
# def add_root(v : Int32) : self
# DP.new(val + cnt, cnt + 1)
# end
# end
# ```
class ReRooting(T, GraphType)
getter graph : GraphType
def initialize(size : Int)
initialize(GraphType.new(size))
end
def initialize(@graph : GraphType)
@dp = Array(Array(T)).new
@result = Array(T).new
end
def initialize(size : Int, edges : Enumerable)
initialize(GraphType.new(size, edges))
end
delegate size, :<<, add_edges, to: @graph
private def dfs(v : Int32, p : Int32) : T
acc = T.new
graph[v].each_with_index do |edge, i|
if edge.to != p
acc += (@dp[v][i] = dfs(edge.to, v))
end
end
acc.add_root(v)
end
private def bfs(v : Int32, p : Int32, dp_p : T) : Nil
graph[v].each_with_index do |edge, i|
@dp[v][i] = dp_p if edge.to == p
end
n = graph[v].size
dp_left = Array.new(n + 1, T.new)
(0...n).each do |i|
dp_left[i + 1] = dp_left[i] + @dp[v][i]
end
dp_right = Array.new(n + 1, T.new)
(0...n).reverse_each do |i|
dp_right[i] = dp_right[i + 1] + @dp[v][i]
end
@result[v] = dp_left.last.add_root(v)
graph[v].each_with_index do |edge, i|
bfs(edge.to, v, (dp_left[i] + dp_right[i + 1]).add_root(v)) if edge.to != p
end
end
def solve : Array(T)
@dp = Array.new(size) { |i| Array.new(@graph[i].size, T.new) }
@result = Array.new(size, T.new)
dfs(0, -1)
bfs(0, -1, T.new)
@result
end
end# require "../graph"
# require "./graph/edge"
struct WeightedEdge(T)
include Comparable(WeightedEdge(T))
property to : Int32, cost : T
def initialize(@to, @cost : T)
end
def <=>(other : WeightedEdge(T))
{cost, to} <=> {other.cost, other.to}
end
def to_s(io) : Nil
io << '(' << to << ", " << cost << ')'
end
def inspect(io) : Nil
io << "->" << to << '(' << cost << ')'
end
end
struct WeightedEdge2(T)
include Comparable(WeightedEdge2(T))
property from : Int32, to : Int32, cost : T
def initialize(@from, @to, @cost : T)
end
def initialize(@from, edge : WeightedEdge(T))
@to, @cost = edge.to, edge.cost
end
def <=>(other : WeightedEdge2(T))
{cost, from, to} <=> {other.cost, other.from, other.to}
end
def reverse : self
WeightedEdge2(T).new(to, from, cost)
end
def sort : self
WeightedEdge2(T).new(*{to, from}.minmax, cost)
end
def to_s(io) : Nil
io << '(' << from << ", " << to << ", " << cost << ')'
end
def inspect(io) : Nil
io << from << "->" << to << '(' << cost << ')'
end
end
struct UnweightedEdge
property to : Int32
def initialize(@to)
end
def initialize(@to, cost)
end
def cost : Int32
1
end
def to_s(io) : Nil
io << to
end
def inspect(io) : Nil
io << "->" << to
end
end
struct UnweightedEdge2
property from : Int32, to : Int32
def initialize(@from, @to)
end
def initialize(@from, @to, cost)
end
def initialize(@from, edge : UnweightedEdge)
@to = edge.to
end
def cost : Int32
1
end
def reverse : self
UnweightedEdge2.new(to, from)
end
def sort : self
UnweightedEdge2.new(*{to, from}.minmax)
end
def to_s(io) : Nil
io << '(' << from << ", " << to << ')'
end
def inspect(io) : Nil
io << from << "->" << to
end
end
module Graph(Edge, Edge2)
include Enumerable(Edge2)
getter graph : Array(Array(Edge))
def initialize(size : Int)
@graph = Array(Array(Edge)).new(size) { [] of Edge }
end
def initialize(size : Int, edges : Enumerable)
initialize(size)
add_edges(edges)
end
# Add *edge*.
abstract def <<(edge : Edge2)
# :ditto:
def <<(edge : Tuple) : self
self << Edge2.new(*edge)
end
def add_edges(edges : Enumerable) : self
edges.each { |edge| self << edge }
self
end
delegate size, :[], to: @graph
# Yields each edge of the graph, ans returns `nil`.
def each(&) : Nil
(0...size).each do |v|
graph[v].each do |edge|
yield Edge2.new(v, edge)
end
end
end
def each_child(vertex : Int, parent, &block) : Nil
graph[vertex].each do |edge|
yield edge if edge.to != parent
end
end
def each_child(vertex : Int, parent)
graph[vertex].each.reject(&.to.== parent)
end
def reverse : self
if self.class.directed?
each_with_object(self.class.new(size)) do |edge, reversed|
reversed << edge.reverse
end
else
dup
end
end
def to_undirected : self
if self.class.directed?
each_with_object(self.class.new(size)) do |edge, graph|
graph << edge << edge.reverse
end
else
dup
end
end
def to_s(io : IO) : Nil
io << '['
join(", ", io) do |edge, io|
edge.inspect io
end
io << ']'
end
def inspect(io : IO) : Nil
io << "[\n"
graph.each do |edges|
io << " " << edges << ",\n"
end
io << ']'
end
end
class DiGraph(T)
include Graph(WeightedEdge(T), WeightedEdge2(T))
def self.weighted?
true
end
def self.directed?
true
end
def initialize(size : Int)
super
end
def initialize(size : Int, edges : Enumerable(WeightedEdge2(T)))
super
end
def initialize(size : Int, edges : Enumerable({Int32, Int32, T}))
super
end
def <<(edge : WeightedEdge2(T)) : self
raise IndexError.new unless 0 <= edge.from < size && 0 <= edge.to < size
@graph[edge.from] << WeightedEdge.new(edge.to, edge.cost)
self
end
end
class UnGraph(T)
include Graph(WeightedEdge(T), WeightedEdge2(T))
def self.weighted?
true
end
def self.directed?
false
end
def initialize(size : Int)
super
end
def initialize(size : Int, edges : Enumerable(WeightedEdge2(T)))
super
end
def initialize(size : Int, edges : Enumerable({Int32, Int32, T}))
super
end
def <<(edge : WeightedEdge2(T)) : self
raise IndexError.new unless 0 <= edge.from < size && 0 <= edge.to < size
@graph[edge.from] << WeightedEdge.new(edge.to, edge.cost)
@graph[edge.to] << WeightedEdge.new(edge.from, edge.cost)
self
end
end
class UnweightedDiGraph
include Graph(UnweightedEdge, UnweightedEdge2)
def self.weighted?
false
end
def self.directed?
true
end
def initialize(size : Int)
super
end
def initialize(size : Int, edges : Enumerable)
super
end
def <<(edge : UnweightedEdge2) : self
raise IndexError.new unless 0 <= edge.from < size && 0 <= edge.to < size
@graph[edge.from] << UnweightedEdge.new(edge.to)
self
end
end
class UnweightedUnGraph
include Graph(UnweightedEdge, UnweightedEdge2)
def self.weighted?
false
end
def self.directed?
false
end
def initialize(size : Int)
super
end
def initialize(size : Int, edges : Enumerable)
super
end
def <<(edge : UnweightedEdge2) : self
raise IndexError.new unless 0 <= edge.from < size && 0 <= edge.to < size
@graph[edge.from] << UnweightedEdge.new(edge.to)
@graph[edge.to] << UnweightedEdge.new(edge.from)
self
end
end
# Example of `T`:
# ```
# struct DP
# getter val : Int64, cnt : Int32
#
# def initialize
# @val, @cnt = 0i64, 0
# end
#
# def initialize(@val, @cnt)
# end
#
# def +(other : self) : self
# DP.new(val + other.val, cnt + other.cnt)
# end
#
# def add_root(v : Int32) : self
# DP.new(val + cnt, cnt + 1)
# end
# end
# ```
class ReRooting(T, GraphType)
getter graph : GraphType
def initialize(size : Int)
initialize(GraphType.new(size))
end
def initialize(@graph : GraphType)
@dp = Array(Array(T)).new
@result = Array(T).new
end
def initialize(size : Int, edges : Enumerable)
initialize(GraphType.new(size, edges))
end
delegate size, :<<, add_edges, to: @graph
private def dfs(v : Int32, p : Int32) : T
acc = T.new
graph[v].each_with_index do |edge, i|
if edge.to != p
acc += (@dp[v][i] = dfs(edge.to, v))
end
end
acc.add_root(v)
end
private def bfs(v : Int32, p : Int32, dp_p : T) : Nil
graph[v].each_with_index do |edge, i|
@dp[v][i] = dp_p if edge.to == p
end
n = graph[v].size
dp_left = Array.new(n + 1, T.new)
(0...n).each do |i|
dp_left[i + 1] = dp_left[i] + @dp[v][i]
end
dp_right = Array.new(n + 1, T.new)
(0...n).reverse_each do |i|
dp_right[i] = dp_right[i + 1] + @dp[v][i]
end
@result[v] = dp_left.last.add_root(v)
graph[v].each_with_index do |edge, i|
bfs(edge.to, v, (dp_left[i] + dp_right[i + 1]).add_root(v)) if edge.to != p
end
end
def solve : Array(T)
@dp = Array.new(size) { |i| Array.new(@graph[i].size, T.new) }
@result = Array.new(size, T.new)
dfs(0, -1)
bfs(0, -1, T.new)
@result
end
end