いものやま。

雑多な知識の寄せ集め

ABC163を復習してみた。

f:id:yamaimo0625:20200420163532j:plain

最近、AtCoderに参戦してる。
これはアルゴリズム実技検定(PAST)を受けてみようと思ったのがキッカケ。

PASTの結果は63点でギリギリ中級だったんだけど、やはり経験値が足りてないなぁという感じ。

なので、ほぼ毎週開かれてるAtCoder Beginner Contest (通称ABC)をちゃんと復習して、経験値を高めていきたいと思う。

先週末に行われたABC163の復習をしてみる。

※なお、ABC163は開始直後にいろいろトラブルがあってレーティングはつかないことになった

A. Circle Pond

円の半径が与えられるので円周の長さを求めるだけ。

r = gets.to_i
puts 2 * Math::PI * r

B. Homework

宿題にかかる日数を合計して、夏休みの日数から引くだけ。
日数が足りない場合は-1を出力せよということなので、そういうときはmaxを使うのが数学のお約束パターン。

n = gets.split.map(&:to_i)[0]
total = gets.split.map(&:to_i).sum()
puts [n-total, -1].max

C. Management

 iさんが jさんの上司の場合(つまり A_j = iの場合)、 jさんは iさんの部下なのだから、 iさんの部下の数を1増やせばいい。
それを Aの先頭から調べていけばいいだけ。

n = gets.to_i
bosses = gets.split.map(&:to_i)

ans = Array.new(n+1,0) # 0はダミー
bosses.each do |boss|
  ans[boss] += 1
end

ans.shift # ダミーを削除
ans.each do |i|
  puts i
end

なお、インデックスを揃えるのが面倒だったので、答えの配列のサイズを一つ大きくして、そのままアクセスできるようにしてる。

このCまでは楽勝だったんだけど、Dで苦戦した・・・

D. Sum of Large Numbers

いろいろ考えて、自分が至ったのは次のような考え:

  1. まず、足す個数が違うときは、和は被らない。
    なので、足す個数がそれぞれ k, k+1, \ldots, n+1のときの和の数を求め、合計すればいい。
  2. 足す個数が i個のとき、一番小さい和は 0 + 1 + 2 + \cdots + (i-2) + (i-1)
    ここから一番最後の項を i-1, i, i+1, \ldots, n-1, nまで大きくしていける。
    これが n - (i-1) +1 = n- i +2通り。
    そしたら最後から2番目の項を i-1, i, i+1, \ldots, n-2, n-1まで大きくしていける。
    これが (n-1)-(i-1)+1=n-i+1通り。
    同様に最後から3番目の項、4番目の項、・・・、最初の項と大きくしていくと、それぞれ n-i+1通り。
    なので、 (n-i+2) + (n-i+1)(i-1)通り。

これをコードに落としたのが、以下:

n, k = gets.split.map(&:to_i)

total = 0
(k..(n+1)).each do |i|
  total += (n-i+2) + (n-i+1)*(i-1)
end

puts total % (10**9+7)

これは時間が足りなくて間に合わなかった。。。

解説読むと、この2.のところがもっとクリアに解決されてる。

まず、一番小さい和は 0+1+\cdots + (i-1)
そして、一番大きな和は (n-i+1)+(n-i+2)+\cdots+n
ここでポイントが、その間の任意の和は好きに作れるということ。
なので、一番小さい和を S_{\min}、一番大きな和を S_{\max}としたとき、

 \begin{align}
S_{\max} - S_{\min} + 1 &= \frac{1}{2} \left((n-i+1)+n \right) i - \frac{1}{2} \left(0 + (i-1)\right) i + 1 \\
&= \frac{1}{2}i(2n-i+1) - \frac{1}{2} i(i-1) + 1 \\
&= i (n-i+1) + 1
\end{align}

通りの和が作れる。

よって、以下のように書ける:

n, k = gets.split.map(&:to_i)

total = 0
(k..(n+1)).each do |i|
  total += i * (n - i + 1) + 1
end

puts total % (10**9+7)

E. Active Infants

活発度が大きい子ほど出来るだけ大きく動かした方がいいので、活発度の大きい子から順に、空きの左端もしくは右端(移動距離の大きくなる方)に移動させていけばいいだけではと思って、その単純な貪欲法を書いてみた。

n = gets.to_i
happy = gets.split.map(&:to_i)
happy.unshift 0 # ダミー

# happyが多い順にインデックスを並び替える
sorted_index = (1..n).sort_by{|i| -happy[i]}

total = 0
first = 1
last = n
sorted_index.each do |i|
  first_dist = (i - first).abs
  last_dist = (last - i).abs
  if first_dist > last_dist
    total += happy[i] * first_dist
    first += 1
  else
    total += happy[i] * last_dist
    last -= 1
  end
end

puts total

ただ、例で試したところ、例1はOKなものの、例2と例3は1足りない・・・
ということで、解けなかった。

どうやら、単に遠い方へ移動させるのだとダメなケースがあるっぽい?
なので、空きの左端に移動させた場合と右端に移動させた場合の両方について考える、動的計画法(DP)が必要とのことだった。

左側にleft人、右側にright人詰めたときの嬉しさの合計の最大値をhappysum[left][right]で表すことにする。
そうすると、その状態になるのは

  1. 左側にleft-1人、右側にright人いて、左側に1人移動させた場合
  2. 左側にleft人、右側にright-1人いて、右側に1人移動させた場合

のいずれか。
そして、直前の状態の嬉しさの最大値はそれぞれ参照できるので、それぞれの場合の嬉しさの合計を計算して、より高い方の値にすればいい、と。

ということで、あとで書き直したのが以下のコード:

n = gets.to_i
happy = gets.split.map(&:to_i)
happy.unshift 0 # ダミー

# happyが多い順にインデックスを並び替える
sorted_index = (1..n).sort_by{|i| -happy[i]}
sorted_index.unshift 0 # ダミー

# テーブルを0で初期化
happysum = (n+1).times.map do |left|  # leftは0, 1, ..., n
  Array.new(n-left, 0) # rightは0, 1, ..., n-left (left+right<=n)
end

# left=0固定
(1..n).each do |right|
  i = sorted_index[right]
  happysum[0][right] = happysum[0][right-1] + happy[i] * (n + 1 - right - i).abs
end

# right=0固定
(1..n).each do |left|
  i = sorted_index[left]
  happysum[left][0] = happysum[left-1][0] + happy[i] * (i - left).abs
end

max = [happysum[0][n], happysum[n][0]].max

# left != 0 かつ right != 0のとき
(1..n).each do |left|
  (1..(n-left)).each do |right|
    i = sorted_index[left+right]
    to_left = happysum[left-1][right] + happy[i] * (i - left).abs
    to_right = happysum[left][right-1] + happy[i] * (n + 1 - right - i).abs
    happysum[left][right] = [to_left, to_right].max
  end
  max = [max, happysum[left][n-left]].max
end

puts max

F. Path pass i

解説を読んだけど、まったく意味不明だったので、ACとってる他の人のコードで理解できそうなものを見つけて解読した。

参考にしたのは以下のコード:

根本的な考え方は、以下のようになる:

まず、サイズ nの木にあるパスの数は、

  1. 始点と終点が異なるものが \frac{1}{2}n(n-1)
  2. 始点と終点が同じものが n

なので、合計 \frac{1}{2}n(n-1) + n = \frac{1}{2}n(n+1)個になる。

そして、「色 cのノードを1回以上通るパスを求めよ」というのは逆に、「色 cを1回も通らないパスの数」を求めて、パスの総数から引けばいい。
(余事象の考え方と同じ)

じゃあ、「色 cを1回も通らないパスの数」をどう求めればいいのか、となるけど、これは色 cのノードを木から取り除いて森にして、森に含まれるパスの数を求めればいい。
その森には色 cのノードがないのだから、色 cを通るパスは存在していない。

これ森に含まれるパスの数は、色 cのノードを取り除いた森に含まれる各木同士は非連結なので、各木にあるパスの数を全部足し合わせれば求まる。

あとは、じゃあ色 cのノードを取り除いた各木のサイズをどうやって求めればいいかだけど、これは深さ優先探索(DFS)をやって、

  1. 今いるノードの色が cの場合、そこで木は切れるので、子を新しい探索のルートとして登録する
  2. 今いるノードの色が cでない場合、木はまだ連結してるので、子のサイズと自分自身を足し合わせる

とすることで求められる。

これをコードにしたのが以下:

class Tree
  def initialize(n, colors, edges)
    @n = n
    @colors = colors
    @adj = @n.times.map{Array.new}
    edges.each do |a, b|
      @adj[a].push b
      @adj[b].push a
    end

    # 計算用
    @cut_color = nil
    @root_queue = []
  end

  def get_answer
    total_path_count = all_path_count(@n)
    (1..@n).map do |color|
      tree_sizes = get_tree_sizes_for_cut_color(color)
      path_count_in_trees = tree_sizes.map{|tree_size| all_path_count(tree_size)}.sum()
      total_path_count - path_count_in_trees
    end
  end

  private

  # colorのノードを取り除いて木を森にしたとき、
  # 森に含まれる木のサイズのリストを得る
  def get_tree_sizes_for_cut_color(color)
    tree_sizes = []
    @cut_color = color
    @root_queue.push [0, -1]
    until @root_queue.empty?
      root, root_parent = @root_queue.shift
      tree_size = get_subtree_size(root, root_parent)
      tree_sizes.push tree_size if tree_size > 0
    end
    tree_sizes
  end

  # node以下の部分木のサイズを返す
  # (nodeが@cut_colorの場合、そこで木は切れるので、
  # 子を新しい木のルートとして登録する)
  def get_subtree_size(node, parent)
    size = 0
    if @colors[node] == @cut_color
      @adj[node].each do |child|
        if child != parent
          @root_queue.push [child, node]
        end
      end
    else
      size = 1 # 自分自身
      @adj[node].each do |child|
        if child != parent
          size += get_subtree_size(child, node)
        end
      end
    end
    size
  end

  # サイズnの木にあるパスの総数
  def all_path_count(n)
    n * (n + 1) / 2
  end
end

n = gets.to_i
colors = gets.split.map(&:to_i)
edges = (n-1).times.map{ gets.split.map{|e| e.to_i - 1} }

tree = Tree.new(n, colors, edges)
tree.get_answer.each do |i|
  puts i
end

これで論理的にはOKなんだけど、走らせるとTLE・・・

そこで、上記は各色について別々にループを回してるけど、これを並列していっぺんに行うことを考える。

具体的には、深さ優先探索をするときに、

  • ノードnode以下の部分木のサイズ
  • 各色について、その色を取り除いたときに、node以下の部分木でnodeから非連結になるノードの数

を返すようにする。
こうすると、nodeから連結してるノードの数は引き算をすれば求まる。

それをコードにしたのが以下:
(※参考にしたコードと等価なコード)

class Tree
  def initialize(n, colors, edges)
    @n = n
    @colors = colors
    @adj = @n.times.map{Array.new}
    edges.each do |a, b|
      @adj[a].push b
      @adj[b].push a
    end

    # 各色でのノードを取り除いて木を森にしたときに、
    # すでに訪問した木に含まれるパスの合計を累積して保持する
    @path_count_in_trees = Hash.new(0)
  end

  def get_answer
    total_path_count = all_path_count(@n)
    size, disconnected_node_sizes = get_subtree_info(0, -1)
    (1..@n).map do |color|
      tree_size = size - disconnected_node_sizes[color]
      @path_count_in_trees[color] += all_path_count(tree_size)
      total_path_count - @path_count_in_trees[color]
    end
  end

  private

  # node以下の部分木について、
  # - サイズ
  # - 各色を取り除いて非連結になるノードの数
  # を返す
  # (サイズ-各色を取り除いて非連結になるノードの数)が、
  # node以下の部分木でnodeから連結してる木のサイズ
  def get_subtree_info(node, parent)
    size = 1
    disconnected_node_sizes = Hash.new(0)

    node_color = @colors[node]
    @adj[node].each do |child|
      if child != parent
        sub_size, sub_disconnected_node_sizes = get_subtree_info(child, node)
        size += sub_size

        # node_colorについて、nodeで木が切断されるので、木のサイズが確定する
        # なので、確定した木に含まれるパスの数を足し込む
        tree_size = sub_size - sub_disconnected_node_sizes[node_color]
        @path_count_in_trees[node_color] += all_path_count(tree_size)

        # 非連結なノード数を足し込む
        # このとき、効率をよくするためにサイズの大きなオブジェクトを再利用する
        if disconnected_node_sizes.size < sub_disconnected_node_sizes.size
          disconnected_node_sizes, sub_disconnected_node_sizes \
            = [sub_disconnected_node_sizes, disconnected_node_sizes]
        end
        sub_disconnected_node_sizes.each do |color, sub_disconnected_node_size|
          disconnected_node_sizes[color] += sub_disconnected_node_size
        end
      end
    end

    # node_colorについて、nodeで木が切断されるので、
    # node以下は親から見て非連結になる
    # なので、非連結なノード数をnode以下の部分木のサイズにする
    disconnected_node_sizes[node_color] = size

    [size, disconnected_node_sizes]
  end

  # サイズnの木にあるパスの総数
  def all_path_count(n)
    n * (n + 1) / 2
  end
end

n = gets.to_i
colors = gets.split.map(&:to_i)
edges = (n-1).times.map{ gets.split.map{|e| e.to_i - 1} }

tree = Tree.new(n, colors, edges)
tree.get_answer.each do |i|
  puts i
end

これでやっとACが取れた。

それにしても、難しい・・・
コード理解するの、めっちゃ時間かかった(^^;
これも慣れればスルッとできるようになるのかなぁ・・・?

今日はここまで!