PythonでEDPC (A~C問題)

競技プログラミング

Python より PyPy

私のような初心者は Python で普通に DP を作ると TLE になる可能性が高いため、

行き詰まったら PyPy で提出したほうが良いかもしれない。

私は今回初めて PyPy を使った。

実際に私は B問題で TLE になり、PyPy に切り替えたら通った。

A - Frog 1

# 入力
N = int(input())
h = list(map(int,input().split()))

# 配列 dp を定義
dp = [0] * N

for i in range(1, N):
    if i == 1:
        dp[i] = abs(h[i] - h[i-1])
    else:
        dp[i] = min(dp[i-1] + abs(h[i] - h[i-1]), dp[i-2] + abs(h[i] - h[i-2]))

# 答え
print(dp[-1])

B - Frog 2

A - Frog1 との違いは カエルが移動できる足場の数が K で示されるという点。

A問題はこの問題における K = 2 の場合と全く同じになる。

だから、K をケアしながら A問題と全く同じように dp テーブルを作ればよい。

# 入力
N, K = map(int, input().split())
h = list(map(int,input().split()))

# 配列 dp を定義
dp = [0] * N

for i in range(1, N):
    # A-Frog1 との違い
    dp[i] = dp[i-1] + abs(h[i] - h[i-1])
    for j in range(1, K+1):
        if i-j >= 0:
            dp[i] = min(dp[i], dp[i-j] + abs(h[i] - h[i-j]))
        # デバッグ用
        # print("i={},j={},{}".format(i,j,dp))

# 答え
print(dp[-1])

C - Vacation

どうやって「2日以上連続で同じ活動を行うことができない」ことを入れ組むかというと、L12 において$ i \neq j $ のときは飛ばすことで成立させている。

計算量

オーダーは $O(3\times 3\times N) = O(9N)$

制約は $1 \le N \le 105 $ より

計算量は $9 \times 105$ なので Python3 でも間に合う

注意

dp を N行3列 と定義しなければならないことに注意

(3行N列にすると、IndexError: list assignment index out of range になってしまう)

解答

# 入力
N = int(input())
abc = [list(map(int,input().split())) for _ in range(N)]

# 配列 dp を定義
dp = [[0] * 3 for _ in range(N)] # 3行N列

# dp[i][j] : i+1 回目までで、最後に行動 j を取ったときの幸福の総和の最大値

# 1 回目
for j in range(3):
    dp[0][j] = abc[0][j]

# 2 ~ N 回目
for i in range(1, N):
    for j in range(3): 
        for k in range(3):
            if j != k:
                dp[i][j] = max(dp[i][j], dp[i-1][k] + abc[i][j])
        
# 答え
ans = max(dp[-1])
print(ans)