This documentation is automatically generated by online-judge-tools/verification-helper

:x: src/graph.cr

Depends on

Required by

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 "./priority_queue.cr"

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
# 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
Back to top page