ABC129 C - Typical Stairs

AtCoder Python3

備忘録

問題

atcoder.jp

回答

import sys
import os

MOD = 10 ** 9 + 7


def main():
    if os.getenv("LOCAL"):
        sys.stdin = open("input.txt", "r")

    N, M = list(map(int, sys.stdin.buffer.readline().split()))
    A = set([int(sys.stdin.buffer.readline().rstrip()) for _ in range(M)])

    ret = [0, 1]
    for i in range(1, N + 1):
        if i in A:
            ret.append(0)
        else:
            ret.append((ret[-1] + ret[-2]) % MOD)

    print(ret[-1] % MOD)


if __name__ == '__main__':
    main()

考え方

AtCoder ABC 129 C - Typical Stairs (300 点) - けんちょんの競プロ精進記録

上記の解説が非常にわかりやすかった。

0段目は０通り、1段目は１通りの移動方法があるため、
初めに[0, 1]の配列を作る。
以降は壊れていない階段であれば、一つ前、あるいは二つ前から登ってくることができるため、
一つ前（ret[-1]）と二つ前(ret[-2]）を足す。

ちなみに、入力を受け取る際に、
A = [int(sys.stdin.buffer.readline().rstrip()) for _ in range(M)]としたらLTEになった。
inで配列内の要素の有無を確認するときはsetオブジェクトのほうが高速っぽい？