verify/graph/lca.test.cr

• View this file on GitHub
• Last update: 2024-07-28 06:26:08+09:00
• Problem: https://judge.yosupo.jp/problem/lca

Depends on

• src/nglib/constants.cr
• src/nglib/graph/lca.cr

Code

# verification-helper: PROBLEM https://judge.yosupo.jp/problem/lca

require "../../src/nglib/graph/lca"

n, q = read_line.split.map &.to_i64
pars = read_line.split.map &.to_i64

graph = Array.new(n) { [] of Int64 }
pars.each_with_index(1) do |par, i|
  graph[i] << par
  graph[par] << i.to_i64
end

lca = NgLib::LCA.new(graph)

ans = String.build { |io|
  q.times do
    u, v = read_line.split.map &.to_i64
    io << lca.ancestor(u, v) << '\n'
  end
}

print ans

# verification-helper: PROBLEM https://judge.yosupo.jp/problem/lca

# require "../../src/nglib/graph/lca"
# require "../constants.cr"
OO = (1_i64 << 62) - (1_i64 << 31)

module NgLib
  # 最近共通祖先を求めるライブラリです。
  class LCA
    alias Graph = Array(Array(Int64))
    getter parent : Array(Array(Int64))
    getter dist : Array(Int64)
    @graph : Graph

    # 木構造グラフ `graph` に対して、`root` を根とする LCA を構築します。
    def initialize(@graph : Graph, root = 0_i64)
      n = graph.size
      k = 1_i64
      while (1_i64 << k) < n
        k += 1
      end

      @parent = Array.new(k) { [-1_i64] * n }
      @dist = [OO] * n
      dfs(root, -1_i64, 0_i64)
      (k - 1).times do |i|
        n.times do |v|
          if @parent[i][v] < 0
            @parent[i + 1][v] = -1_i64
          else
            @parent[i + 1][v] = @parent[i][@parent[i][v]]
          end
        end
      end
    end

    # 頂点 $u$ と 頂点 $v$ の最近共通祖先を返します。
    def ancestor(u : Int, v : Int) : Int64
      if @dist[u] < @dist[v]
        u, v = v, u
      end
      n = @parent.size
      n.times do |k|
        u = @parent[k][u] if (dist[u] - dist[v]).bit(k) == 1
      end
      return u if u == v
      (n - 1).downto(0) do |k|
        if @parent[k][u] != @parent[k][v]
          u, v = @parent[k][u], @parent[k][v]
        end
      end
      @parent[0][u]
    end

    # 頂点 $u$ と頂点 $v$ の距離を返します。
    def distance_between(u : Int, v : Int) : Int64
      dist[u] + dist[v] - dist[ancestor(u, v)] * 2
    end

    # 頂点 $u$ から頂点 $v$ までのパスに頂点 $a$ が含まれているか返します。
    def on_path?(u : Int, v : Int, a : Int) : Bool
      distance_between(u, a) + distance_between(a, v) == distance_between(u, v)
    end

    private def dfs(root : Int64, par : Int64, d : Int64)
      @parent[0][root] = par
      @dist[root] = d
      @graph[root].each do |child|
        next if child == par
        dfs(child, root, d + 1)
      end
    end
  end
end

n, q = read_line.split.map &.to_i64
pars = read_line.split.map &.to_i64

graph = Array.new(n) { [] of Int64 }
pars.each_with_index(1) do |par, i|
  graph[i] << par
  graph[par] << i.to_i64
end

lca = NgLib::LCA.new(graph)

ans = String.build { |io|
  q.times do
    u, v = read_line.split.map &.to_i64
    io << lca.ancestor(u, v) << '\n'
  end
}

print ans

Back to top page