問題
要約
-
長さNの整数列A = (A1, ..., AN)を考える。
-
以下の条件を満たす数列の個数を求める
a. 各要素Aiは1以上M以下の整数(1 ≤ i ≤ N)。
b. 数列の要素の合計がK以下(A1 + A2 + ... + AN ≤ K)。 -
答えが非常に大きくなる可能性があるため、998244353で割った余りを求める。
- N: 数列の長さ
- M: 各要素の最大値
- K: 数列の要素の合計の上限
- Ai: 数列の各要素(1 ≤ Ai ≤ M)
既存投稿一覧ページへのリンク
アプローチ
動的計画法を用いて、条件を満たす数列の個数を計算する。
二次元のDPテーブルを使用し、数列の長さと要素の合計を考慮しながら、可能な組み合わせの数を累積的に計算する。
解法手順
-
入力値N, M, Kを受け取る。
-
(N+1) × (K+1)の二次元配列dpを作成し、初期値を0で埋める。
-
dp[0][0]を1に設定する(長さ0で合計0の数列は1通り)。
-
以下のループを実行する:
a. 数列の長さiを0からN-1まで繰り返す。
b. 現在の合計jを0からK-1まで繰り返す。
c. 追加する要素lを1からMまで繰り返す。
d. j+l ≤ Kの場合、dp[i+1][j+l] += dp[i][j]を実行する。 -
dp[N][1]からdp[N][K]までの値を合計し、答えとする。
-
答えをmod(998244353)で割った余りを計算する。
-
最終的な答えを出力する。
ACコード
# モジュロ演算の基数
mod = 998244353
def io_func():
# 入力値N, M, Kを受け取る
return map(int, input().split())
def solve(n, m, k):
# (N+1) × (K+1)の二次元配列dpを作成し、初期値を0で埋める
dp = [[0 for _ in range(k+1)] for _ in range(n+1)]
# dp[0][0]を1に設定する(長さ0で合計0の数列は1通り)
dp[0][0] = 1
# 数列の長さiを0からN-1まで繰り返す
for i in range(n):
# 現在の合計jを0からK-1まで繰り返す
for j in range(k):
# 追加する要素lを1からMまで繰り返す
for l in range(1, m+1):
# j+l ≤ Kの場合、dp[i+1][j+l] += dp[i][j]を実行する
if j + l <= k:
dp[i+1][j+l] += dp[i][j]
# dp[N][1]からdp[N][K]までの値を合計し、答えとする
ans = sum(dp[n][1:k+1])
# 答えをmod(998244353)で割った余りを計算する
ans %= mod
# 最終的な答えを返す
return ans
if __name__=="__main__":
# メイン処理
n, m, k = io_func()
result = solve(n, m, k)
print(result)
###
# n: 数列の長さ
# m: 各要素の最大値
# k: 数列の要素の合計の上限
# dp: 動的計画法のテーブル。dp[i][j]は長さiで合計jの数列の個数
# ans: 条件を満たす数列の総数
# 1. io_func関数: 標準入力から値を読み取る
# 2. solve関数: 動的計画法を用いて問題を解く
# - dpテーブルの初期化
# - 三重ループでdpテーブルを埋める
# - 最終的な答えを計算し、モジュロ演算を適用
# 3. メイン処理: io_func関数で入力を受け取り、solve関数で結果を計算し、出力する
思考過程
数列の長さと要素の合計を考慮しながら、可能な組み合わせの数を累積的に計算する。
計算量の概算
(N * K * M)
その他周辺知識
2次元リストの初期化
dp = [[0 for _ in range(k+1)] for _ in range(n+1)]
これは2次元のリストを作成している。
- 外側のリスト内包表記:
[... for _ in range(n+1)]
で、n+1個の要素を持つリストを作成。 - 内側のリスト内包表記:
[0 for _ in range(k+1)]
で、各要素がk+1個の0を持つリストになる。
結果として、(n+1) × (k+1)の大きさの2次元リストが作成され、すべての要素が0で初期化される。
動的計画法(漸化式)
for i in range(n):
for j in range(k):
for l in range(1, m+1):
if j + l <= k:
dp[i+1][j+l] += dp[i][j]
dp[i+1][j+l] += dp[i][j]という漸化式を用いて、より小さな部分問題の解から大きな問題の解を構築している。
ここで、
- i は現在の数列の長さ
- j は現在の数列の要素の合計
- l は新たに追加する要素の値
を表している。
例えば、dp[3][10](長さ3で合計10の数列の個数)を計算する際、dp[2][7]、dp[2][8]、dp[2][9]などの値を利用する。
これは「長さ2で合計7の数列に3を追加する」「長さ2で合計8の数列に2を追加する」などの操作に対応している。
累積和
ans = sum(dp[n][1:k+1])
最終的な答えを求める際、dp[n][1]からdp[n][k]までの値を合計している。
エッセンス
def count_sequences(n, m, k, mod=998244353):
"""
条件を満たす数列の組み合わせの数を計算する関数
Parameters:
n (int): 数列の長さ
m (int): 各要素の最大値
k (int): 数列の要素の合計の上限
mod (int): モジュロ演算の基数 (デフォルト: 998244353)
Returns:
int: 条件を満たす数列の組み合わせの数(モジュロ mod)
"""
# (N+1) × (K+1)の二次元配列dpを作成し、初期値を0で埋める
dp = [[0 for _ in range(k+1)] for _ in range(n+1)]
# dp[0][0]を1に設定する(長さ0で合計0の数列は1通り)
dp[0][0] = 1
# 数列の長さiを0からN-1まで繰り返す
for i in range(n):
# 現在の合計jを0からK-1まで繰り返す
for j in range(k):
# 追加する要素lを1からMまで繰り返す
for l in range(1, m+1):
# j+l ≤ Kの場合、dp[i+1][j+l] += dp[i][j]を実行する
if j + l <= k:
input(dp)
dp[i+1][j+l] = (dp[i+1][j+l] + dp[i][j]) % mod
input(dp)
# dp[N][1]からdp[N][K]までの値を合計し、答えとする
ans = sum(dp[n][1:k+1]) % mod
return ans
if __name__=="__main__":
# 使用例
n = 3 # 数列の長さ
m = 4 # 各要素の最大値
k = 7 # 数列の要素の合計の上限
result = count_sequences(n, m, k)
print(f"条件を満たす数列の組み合わせの数: {result}")
# カスタムのモジュロ基数を使用する場合
custom_mod = 1000000007
result_custom_mod = count_sequences(n, m, k, mod=custom_mod)
print(f"カスタムモジュロを使用した結果: {result_custom_mod}")
動的計画法の補足
問題
が求めたいのは、条件を満たす数列の個数である。
この条件とは:
- 数列の長さがN
- 各要素が1以上M以下
- 数列の要素の合計がK以下
動的計画法の考え方
- 長さNの数列の個数を求めるためには、長さN-1の数列の個数が分かっていればよい。
- そのため、DPテーブルをdp[i][j]のように作り、iを数列の長さ、jを数列の要素の合計とする。
- dp[i][j]は、長さiで合計jの数列の個数を表す。
DPテーブルの構築
- 初期状態: dp[0][0] = 1 (長さ0で合計0の数列は1通り)
- dp[i+1][j+l] += dp[i][j] (長さiで合計jの数列に、要素lを追加)
計算の流れ:
- 数列の長さiを0からN-1まで増やす。
- 各長さiに対して、現在の合計jを0からK-1まで考える。
- 各(i, j)の組み合わせに対して、新たに追加する要素lを1からMまで試す。
- j+l ≤ Kの場合、dp[i+1][j+l]にdp[i][j]を加算する。
最終的な答え:
- dp[N][1]からdp[N][K]までの値を合計すると、条件を満たす数列の総数が得られる。
、例えばN=3, M=2, K=4の場合
dp[0]: [1, 0, 0, 0, 0]
dp[1]: [0, 1, 1, 0, 0]
dp[2]: [0, 0, 1, 2, 1]
dp[3]: [0, 0, 0, 1, 3]
最終的な答えは、dp[3][1] + dp[3][2] + dp[3][3] + dp[3][4] = 0 + 0 + 1 + 3 = 4 となる。