はじめに

ある日のこと。
とあるアルゴリズム問題に競技プログラマー御用達のプログラミング言語である C++ でチャレンジしていたところ、こんなことを思いました。

「これ、Python で解きたくない……？？？」

そう、普段 Python ばかり書いている自分に C++ は難しすぎたのです。
アルゴリズムを考えている時間より C++ で綺麗に実装する方法を考えている時間の方が長いという有様でした。

しかし、Python の弱点は C++ より処理が遅いということ（forループの遅さなどが有名ですね）。
高速化できるよう、工夫して解く練習をする必要がありました。

というわけで C++ に反旗を翻し、Python で Atcoder の Frog 1, Frog 2 を解いていきます。
また、Atcoder には Pypy（高速で動く Python みたいなもの）が存在しますが、Python の高速化についても考えるため、今回は使用しないものとします。

この記事の内容

この記事で扱うこと

C++(GCC 9.2.1) で解いた場合のコードの掲載
Python(3.8.2) で Frog 1, Frog 2 を解く方法
DPの構築方法の軽い解説

この記事で扱わないこと

C++, Python などの文法の解説
Pypyは使わない（普通に Atcoder するなら使ったほうが良いです）
DPの構築方法の細かい解説

A - Frog1

問題：A - Frog 1
種別：最小化問題

最初はコストの最小化を行う問題ですね。

足場を順番に見ていき、一つ前と二つ前の足場のどちらから飛んでくる方がコストが小さいのかを確認することを繰り返すと良さそうです。
また、i=0とi=1の場合には一つ前、二つ前は存在しませんから、別途考える必要がありそうですね。

つまり以下のようになります。

数式

dp_0 = 0

dp_1 = |h_1 - h_0|

dp_i = min(dp_{i-1}+|h_i - h_{i-1}|, dp_{i-2} + |h_i - h_{i-2}|)

C++

これを C++ に落とし込むと以下のようになりますね。

#include <iostream>
#include <vector>

using namespace std;
typedef long long ll;
const ll INF =  1LL << 60;

int N;
vector<int> h;

vector<ll> dp;

int solve (void) {

  // テーブル定義
  dp.assign(N, INF);
  
  // 初期化
  dp[0] = 0;
  dp[1] = abs(h[1] - h[0]);
  
  for (int i = 2; i < N; i++) {
    // 一つ前と二つ前の飛んできた後の合計コストが低い方を代入
    dp[i] = min(
              dp[i - 1] + abs(h[i] - h[i - 1]),
              dp[i - 2] + abs(h[i] - h[i - 2])
            );
  }

  // 最後の足場の最小コストを返す
  return dp[N - 1];
}

int main (void) {

  cin >> N;
  h.assign(N, 0);
  for (int i = 0; i < N; i++) {
    cin >> h[i];
  }
  
  int result = solve();
  cout << result << endl;

  return 0;
}

Python

この程度のDPであればPythonで普通に書いても通ります。

from typing import List

def solve (N: int, h: List[int]) -> int:

    # テーブル定義
    dp = [float('inf')] * N
    
    # 初期化
    dp[0] = 0
    
    for i in range(2, N):
        # 一つ前と二つ前の飛んできた後の合計コストが低い方を代入
        dp[i] = min([x + abs(h[i] - h[max(0, i - 2) + j]) for j, x in enumerate(dp[max(0, i - 2): i])])

    return dp[N-1]

def main () -> None:

    N = int(input())
    h = [int(x) for x in input().split()]

    result = solve(N, h)
    print(result)

    return

if __name__ == '__main__':
    main()

結果

それぞれのスコアは以下のような感じです。

この時点で約5.5倍の違いがあります。

B - Frog2

問題：B - Frog 2
種別：最小化問題

二問目は一問目の発展問題ですね。
二つ前までの足場だけでなく、K個前の足場まで参照しなければなりません。

つまり以下のようになります。

dp_0 = 0

dp_i = min(dp_{i-1} + |h_i - h_{i-1}|, dp_{i-2} + |h_i - h_{i-2}|, ... , dp_{i-K} + |h_i - h_{i-K}|)

C++

これを C++ に落とし込むと以下のようになりますね。

#include <iostream>
#include <vector>

using namespace std;
typedef long long ll;
const ll INF =  1LL << 60;

int N;
int K;
vector<int> h;

vector<ll> dp;

int solve (void) {

  // テーブル定義
  dp.assign(N, INF);

  // 初期化
  dp[0] = 0;

  for (int i = 0; i < N; i++) {
    // K個前までの足場から飛んできた後の合計コストが低いものを代入
    for (int j = 1; i + j < N && j <= K; j++) {
      dp[i + j] = min(
                    dp[i + j],
                    dp[i] + abs(h[i + j] - h[i])
                  );
    }
  }

  // 最後の足場の最小コストを返す
  return dp[N - 1];
}

int main (void) {

  cin >> N;
  cin >> K;
  h.assign(N, 0);
  for (int i = 0; i < N; i++) {
    cin >> h[i];
  }

  int result = solve();
  cout << result << endl;

  return 0;
}

Python

そのままPythonに書き直してみましょう。

from typing import List
 
def solve (N: int, K: int, h: List[int]) -> int:

    # テーブル定義
    dp = [float('inf')] * N
    
    # 初期化
    dp[0] = 0
    
    for i in range(1, N):
        # K個前までの足場から飛んできた後の合計コストが低いものを代入
        dp[i] = min([x + abs(h[i] - h[max(j, i - K + j)]) for j, x in enumerate(dp[max(0, i - K): i])])
 
    return dp[N-1]
 
def main () -> None:

    N, K = [int(x) for x in input().split()]
    h = [int(x) for x in input().split()]
 
    result = solve(N, K, h)
    print(result)
 
    return
 
if __name__ == '__main__':
    main()

というわけでこれを嬉々として提出しに行くわけですが……

まぁ、そうなりますよね。
美しい黄色の TLE(Time Limit Exceeded): 時間切れ を見ることができました。

気を取り直して高速化していきましょう。

標準入力を高速化する

実はPythonの標準入力はあんまり早くありません。
しかしinput()の使いやすさは捨てがたいので、ここはinput()関数を上書きしてしまいましょう。

import sys
input = sys.stdin.readline

これで高速な入力が使えるようになりました。

ループを排除する

恐らく最もクリティカルにTLEの原因となっているであろう箇所は、以下のdpテーブルの更新部分ですね。

for i in range(1, N):
    # K個前までの足場から飛んできた後の合計コストが低いものを代入
    dp[i] = min([x + abs(h[i] - h[max(j, i - K + j)]) for j, x in enumerate(dp[max(0, i - K): i])])

リスト内包表記になっている為、二重ループになっていますね。
内側のループを排除する方法を考えてみましょう。

ここでリスト内包表記のforを使っているのは、abs()関数は入力としてリストなどを受け付け、纏めて絶対値化することができないからです。
つまり、纏めて絶対値にする方法があれば内側のforを排除できそうです。

そこで活躍するのが我らがnumpyモジュールのnumpy.abs()です。これは複数の値を纏めて絶対値に変換することができます。
以下のようにしてみましょう。

import numpy as np

for i in range(1, N):
    # K個前までの足場から飛んできた後の合計コストが低いものを代入
    # numpy を使うと纏めて abs できる
    dp[i] = min(dp[max(0, i-K): i] + np.abs(h[i] - h[max(0, i-K): i]))

改善後

というわけでこんな感じになりました。

import sys
from typing import List

import numpy as np

def solve (N: int, K: int, h: List[int]) -> int:

    # テーブル定義
    dp = np.zeros(N, dtype=int)
    h = np.array(h)

    for i in range(1, N):
        # K個前までの足場から飛んできた後の合計コストが低いものを代入
        # numpy を使うと纏めて abs できる
        dp[i] = min(dp[max(0, i-K): i] + np.abs(h[i] - h[max(0, i-K): i]))

    return dp[N-1]

def main () -> None:

    # 入力関数を sys.stdin.readline に変更
    input = sys.stdin.readline

    N, K = [int(x) for x in input().split()]
    h = [int(x) for x in input().split()]

    result = solve(N, K, h)
    print(result)

    return

if __name__ == '__main__':
    main()

結果

それぞれのスコアは以下のような感じです。

約29倍もの違いがありますね。
このように、PythonでAtcoderなどの競技プログラミングに挑戦する場合、かなり気を使ったコーディングが求められます。

おわりに

やっぱりC++って圧倒的に早いですね……。普通にAtcoderやるなら素直にC++を使ったほうが良いと思います。

おまけ

Frog 2 の改善していない方のコードをPyPy3(7.3.0)に投げ込んでみました。

すごく、はやい。