C - management(AtCoder Beginner Contest 163復習)

4月19日、AtCoder Beginner Contest 163Rubyで参加しました。実行時間オーバーでだめだったC問題に再チャレンジします。

masuyama13.hatenablog.com

C - management

C - management

問題文

問題文
N 人の社員からなる会社があり、各社員には 1,...,N の社員番号が割り当てられています。
社員番号 1 の社員以外の全ての社員には、自分より社員番号が小さい直属の上司がちょうど 1 人います。
X さんが Y さんの直属の上司であるとき、Y さんは X さんの直属の部下であるといいます。
社員番号 i の社員の直属の上司の社員番号が Ai であることが与えられます。各社員について直属の部下が何人いるか求めてください。

制約
2 ≤ N ≤2×105, 1 ≤ Ai < i

入力
入力は以下の形式で標準入力から与えられる。
N
A2...AN

出力
社員番号 1,2,...,N のそれぞれの社員について、直属の部下が何人いるか、改行区切りで出力せよ。

(実行時間制限: 2 sec / メモリ制限: 1024 MB)

提出したコード

n = gets.to_i
numbers = gets.split.map(&:to_i)
(1..n).each do |x|
    puts numbers.count(x)
end

結果:TLE(実行時間超過) 実行時間:2206 ms

考えたことを言語化してみる

最初問題を読んだとき、全然わからなくて諦めそうになりました。それでD問題を見たりしましたが、当然さらに難しい…。時間もまだあったので、とりあえずやってみることに。

文章をよくよく読むと、部下の人数を上司の社員番号順に出力すればいいということがわかりました。上司の社員番号は与えられるので、それを数えて順番に出力するだけなら意外と簡単!

1行目:全体の社員数nを取得。

2行目:標準入力の2行目が、各社員の上司の社員番号です。全体をgetsで取得してsplitでスペース区切りで配列にし、mapで各要素の文字列を整数に変換してその配列にnumbersという名前をつける(mapの解説はこちら)。

3行目以降(each文):countで配列numbersの中に1(x)がいくつあるか数え、その数を出力。その数というのは、社員番号1の人が上司になっている数、つまり社員番号1の人の部下の人数。次に2(x)がいくつあるか数え、その数を出力・・・。これをn回繰り返す。

改善点など

実行時間超過になったのが初めてで原因がわからなかったので、仮説を立てて調べてみました。

(仮説1)putsを繰り返し処理したため?

ということで、出力したいものをいったん配列(array)に入れておいて、繰り返し処理が終わってから最後にその配列を出力するように変更してみました。

n = gets.to_i
numbers = gets.split.map(&:to_i)
array = []
(1..n).each do |x|
  array << numbers.count(x)
end
puts array

結果:TLE(実行時間超過) 実行時間:2206 ms

見当違いだったようです。最初のコードよりも悪くなった気がします。

(仮説2)countが遅い?

正直、遅い原因に関して見当もつかないのですが、eachはよく使っていますし、これを使うなと言われたらコードが書けません。となると、今回初めて使ったcountが遅いのではないか、ということに。とはいえどう書けばいいかわからないので、他の参加者の皆様のコードを参考にさせていただきました。

改善後のコード
n = gets.to_i
numbers = gets.split.map(&:to_i)
array = Array.new(n, 0)
numbers.each do |x|
  array[x - 1] += 1
end
puts array

結果:AC(正解) 実行時間:187 ms

1・2行目は前と同じです。

3行目:新しい配列を作るArray.newには引数を与えることができます。1つ目の引数で要素の数を指定、2つ目の引数で要素の初期値を指定(詳細はリファレンス参照)。今回の例では、n個(全社員の人数分)要素を持ち、各要素の値が0の配列を作ります。そしてその配列にarrayという名前をつけます。このarrayが、最後に出力する答えになっていきます。

4行目:配列numbersを繰り返し処理します。 numbersの1つ目の要素がxならば、array[x - 1]1を足します。 たとえば要素の値が3だったとすると、上司の社員番号が3ということなので、array[2]にプラス1をします。 array[x - 1]と1を引くのは、配列の添字(インデックス)が0から始まるからですね。社員番号は1から始まるので、そこを揃えるためです。 これを2つ目の要素、3つ目の要素・・・と繰り返す。

5行目:最後に配列arrayを出力。

さて、実行速度の話に戻ります。下記記事のコメントによると、countは内部でeachを使っているようです。

Array や Hash には Enumerable が include されているから count が使えるのですね。
で,この Enumerable#count は each を介して数えるから遅いです。
length、size、count メソッドの違いまとめ【Ruby】 - Qiita

つまり、最初のコードだと、each文の中でさらにeachを使っているようなものですね。それは遅そうです。これからは速度にも気をつけたいと思います。

ちなみに、下のようにもっと短く書くこともできるみたいです。いろいろな書き方があって勉強になります!

array = [0] * gets.to_i
gets.split.map(&:to_i).each { |x| array[x - 1] += 1 }
puts array

(他参加者のコードを参考に一部改変。)