【ruby】第２回全国プログラミング王決定戦予選で出てきた謎問題

全国プログラミング王決定戦予選とは

ヤサゴリがいつもやってるAtCorderで行われるコンテストの1つ。
産経主催のコンテストで通常のAtCorderと少しテイストが異なる。
初チャレンジで見事どハマりしたから、忘れないようにメモ。

連結無向グラフとは

躓いたのはこの問題。

頂点に 1 から N の番号が付けられた N 頂点からなる木

普通に木を想像してたけど、全然違った。問題に出ている入力例を絵にするとこんな感じ。 入力例１だと分かりづらいけど、入力例２でだいぶわかる。

7
0 3 2 1 2 2 1

つまり絵の赤字と入力された数字が一致している。

１文字追加して考える

さらに１文字追加するとこんな感じ。

8
0 3 2 1 2 2 1 3

数学的な答えは絵に書いてある2³ * 3² = 72通りになる。

解き方

まずは値を入力させる。

n = gets.to_i
ds = gets.split.map(&:to_i)
mod = 998244353  # これで割った余りを出すのでとりあえず定義

上で考えた通りds内で各数字の個数をまとめていきたいのでgroup_byが使えそう。

group_byメソッド

同じ言葉・数字同士でまとめることができる。

ary = [0, 3, 2, 1, 2, 2, 1]
ary.sort.group_by(&:itself)
=> {0=>[0], 1=>[1, 1], 2=>[2, 2, 2], 3=>[3]}
ary.sort.group_by(&:itself).map{|k, v| [k, v.count]}.to_h
=> {0=>1, 1=>2, 2=>3, 3=>1}

group_byはハッシュ化されるので今回は使い勝手が悪そう。 なので、素直に配列化させる。

cnt = Array.new(ds.max+1, 0)
(0..n-1).each do |k|
  if (k == 0 && ds[k] !=0) || (k != 0 && ds[k] == 0)
    puts 0  # この条件の時は木が作れないから無条件で0を出力して終了
    exit
  end
  cnt[ds[k]] += 1
end

# ds = [0, 3, 2, 1, 2, 2, 1]
# => cnt = [1, 2, 3, 1]

0は最初の１回しか使わないので、この後の計算では不要なので消しておく。この時に使うのはshiftメソッド。

shiftメソッド

配列の最初から数えてn個の要素を削除する。

ary = ["Shiba", "Pochi", "Koko", "Maru"]
ary.shisft
puts ary
=> ["Pochi", "Koko", "Maru"]
ary = ["Shiba", "Pochi", "Koko", "Maru"]
ary.shift(2)
puts ary
=>  ["Koko", "Maru"]

解き方 続き

ここで改めてcase3'で考えてみる。 青・緑・黄色で囲んだ数字がそれぞれ一致している。これをコードに反映する。

cnt.shift
ans = 1
(0..cnt.length-2).each do |k|
  ans = ans * (cnt[k] ** cnt[k + 1])
end
ans %= mod  # modで割った余りを出力するので忘れずに
puts ans

これでOK。

まとめ

全てまとめるとこうなる

n = gets.to_i
ds = gets.split.map(&:to_i)
mod = 998244353
cnt = Array.new(ds.max+1, 0)
(0..n-1).each do |k|
  if (k == 0 && ds[k] !=0) || (k != 0 && ds[k] == 0)
    puts 0
    exit
  end
  cnt[ds[k]] += 1
end
cnt.shift
ans = 1
(0..cnt.length-2).each do |k|
  ans = ans * (cnt[k] ** cnt[k + 1])
end
ans %= mod
puts ans

これは紙とペンがないと辛い問題だった…配点が600点というのも納得。（BeginnerContestならFランク問題…！）