ABC179 D - Leaping Tak - ペンギンかっこいい

備忘録

問題

回答

import sys
import os
import math
import bisect
import itertools
import collections
import heapq
import queue
import array

# 時々使う
# from scipy.sparse.csgraph import csgraph_from_dense, floyd_warshall
# from decimal import Decimal
# from collections import defaultdict, deque

# 再帰の制限設定
sys.setrecursionlimit(10000000)


def ii(): return int(sys.stdin.buffer.readline().rstrip())
def il(): return list(map(int, sys.stdin.buffer.readline().split()))
def fl(): return list(map(float, sys.stdin.buffer.readline().split()))
def iln(n): return [int(sys.stdin.buffer.readline().rstrip())
                    for _ in range(n)]


def iss(): return sys.stdin.buffer.readline().decode().rstrip()
def sl(): return list(map(str, sys.stdin.buffer.readline().decode().split()))
def isn(n): return [sys.stdin.buffer.readline().decode().rstrip()
                    for _ in range(n)]


def lcm(x, y): return (x * y) // math.gcd(x, y)


# MOD = 10 ** 9 + 7
MOD = 998244353
INF = float('inf')


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

    N, K = il()
    D = [il() for _ in range(K)]

    dp = [0] * (N)
    dp[0] = 1
    acc = 0
    for i in range(1, N):
        for l, r in D:
            if i - l >= 0:
                acc += dp[i-l]
                acc %= MOD
            if i - r - 1 >= 0:
                acc -= dp[i-r-1]
                acc %= MOD
        dp[i] = acc
    print(dp[-1])


if __name__ == '__main__':
    main()

考え方

累積和を用いた動的計画法で回答を得る。

dp[i]を位置iに到達することが出来る移動方法の総数とし、
動的計画法でdp[N]を求める。

移動方法には与えられたK種類の区間で移動することができ、
例えば区間[1, 3]を使用して位置4に到達することが出来るのは、
位置1, 2, 3の３つから位置4に移動することができる。

この位置4に到達することができる移動方法の総数を求めるために、累積和を用いる。

例えば、与えられる区間が[1, 3]のみのとき、dp[i]の遷移は以下の通りになる。

# dp[0]は初期位置
# 到達することが出来る総数はスタート時点の１のみなので、
# dp[0] = 1で初期化を行う
dp = [1, 0, 0, 0, 0, ... , 0]

# dp[0]から区間[1, 3]で移動を行った場合
# dp[1]にはdp[0] + 1の計1通りの移動方法で到達可能
dp = [1, 1, 0, 0, 0, ... , 0]

# dp[2]にはdp[0] + 2, dp[1] + 1の計2通りの移動方法で到達可能
dp = [1, 1, 2, 0, 0, ... , 0]

# dp[3]にはdp[0] + 3, dp[1] + 2, dp[2] + 1の計4通りの移動方法で到達可能
dp = [1, 1, 2, 4, 0, ... , 0]

dp[4]への移動方法を求めるとき、
区間[1, 3]を使用すると、dp[1], dp[2], dp[3]からdp[4]に移動可能。
つまり、dp[1], dp[2], dp[3]への移動可能な総数がdp[4]の値となる。
(dp[4] = dp[1] + dp[2] + dp[3])

これは区間l, rの和と考えることができ、
指定区間の和は累積和で高速に求めることができる。

# dp[4] = dp[1] + dp[2] + dp[3]の計7通りの移動方法で到達可能
dp = [1, 1, 2, 4, 7, ... , 0]

# これを高速に求めるために、
# 予めdp配列の合算した値を用意しておき、累積和の要領で
# dp[4]の値を求める
acc = dp[0] + dp[1] + dp[2] + dp[3]
dp[4] = acc - dp[0] # O(1)でdp[4]を求めることができる

あとは配列の範囲外を指定しないよう気を付けて、
累積和を計算することでdp[N]を求めることができる。