Ruby と Perl と Java で解く AtCoder ABC 129 C (後編) 動的計画法

はじめに

AtCoder Problems の Recommendation を利用して、過去の問題を解いています。
AtCoder さん、AtCoder Problems さん、ありがとうございます。

今回のお題

AtCoder Beginner Contest 129 C - Typical Stairs
Difficulty: 795

*前編*のテーマは、フィボナッチ
後編のテーマは、動的計画法

Ruby

階段数の組合せがフィボナッチ数に等しかったので、前編の解法が使えましたが、もっと汎用的に解けないものでしょうか。

そこで、オフィシャルの editorial にあります動的計画法の登場になります。

ruby.rb

n, m = gets.split.map(&:to_i)
MOD = 1_000_000_007
dp = Array.new(n + 2, 0)
a = []
dp[0] = 1
m.times do |i|
  a[i] = gets.to_i
  dp[a[i]] = -1
end
n.times do |i|
  next if dp[i] == -1
  if dp[i + 1] != -1
    dp[i + 1] += dp[i]
    dp[i + 1] %= MOD
  end
  if dp[i + 2] != -1
    dp[i + 2] += dp[i]
    dp[i + 2] %= MOD
  end
end
puts dp[n]

array.rb

  oks[a]=false;  # => 壊れた床の配列

  dp[a[i]] = -1  # => 壊れた床を -1

オフィシャルの editorial は、壊れた床用の配列を用意されていますが、ここでは、DP 配列の中に組み込みます。下の図のウォッチ式で、-1 の部分が壊れた床です。

VSCode 等でデバッグ実行した際、変数やウォッチ式に DP 配列の中が表示されるので、数値の変化や壊れた床をスキップする様子を確認することができます。
Ruby AtCoder向けVSCode設定
PerlのAtCoder向けVSCode設定(tasks.json)

fib版.rb

  if a[i] - a[i - 1] == 1
    puts "0"
    exit
  end

フィボナッチ版のソースにあった、壊れた床が連続しているかどうかを判定している箇所が、動的計画法版のソースにはありません。
ありませんが、壊れた床が連続した場合、ちゃんと 0 を返しています。

また、例えば 3 段上ることができるといった条件変更にも、容易に対応可能です。

動的計画法、凄い。

Perl

perl.pl

chomp (my ($n, $m) = split / /, <STDIN>);
my $MOD = 1e9 + 7;
my (@a, @dp);
$dp[0] = 1;
for my $i (1..$m) {
  chomp (my $in = <STDIN>);
  $a[$i] = $in;
  $dp[$a[$i]] = -1;
}
for my $i (0..$n - 1) {
  next if $dp[$i] == -1;
  if ($dp[$i + 1] != -1) {
    $dp[$i + 1] += $dp[$i];
    $dp[$i + 1] %= $MOD;
  }
  if ($dp[$i + 2] != -1) {
    $dp[$i + 2] += $dp[$i];
    $dp[$i + 2] %= $MOD;
  }
}
print $dp[$n] + 0, "\n";

Perl は、未定義値が計算に使用された場合、0 として扱われることを利用しております。

Java

java.java

import java.util.*;

class Main {
    public static void main(String[] args) {
        Scanner sc = new Scanner(System.in);
        int n = Integer.parseInt(sc.next());
        int m = Integer.parseInt(sc.next());
        int MOD = 1_000_000_007;
        int a[] = new int[m];
        int dp[] = new int[n + 2];
        dp[0] = 1;
        for (int i = 0; i < m; i++) {
            a[i] = Integer.parseInt(sc.next());
            dp[a[i]] = -1;
        }
        sc.close();
        for (int i = 0; i < n; i++) {
            if (dp[i] == -1) {
                continue;
            }
            if (dp[i + 1] != -1) {
                dp[i + 1] += dp[i];
                dp[i + 1] %= MOD;
            }
            if (dp[i + 2] != -1) {
                dp[i + 2] += dp[i];
                dp[i + 2] %= MOD;
            }
        }
        System.out.print(dp[n]);
    }
}

今回は足し算なので、long を使用しなくても AC です。

	Ruby	Perl	Java
コード長	359 Byte	436 Byte	902 Byte
実行時間	64 ms	83 ms	347 ms
メモリ	3912 KB	12288 KB	44684 KB

まとめ

• ABC 129 C を苦労して解いた
• Ruby に詳しくなった
• Perl に詳しくなった
• Java に詳しくなった
• 動的計画法に詳しくなった

参照したサイト
ヤマカサの競技プログラミング [AtCoder] ABC 129 C – Typical Stairs