問題
既存投稿一覧ページへのリンク
階乗と逆元の前計算
sample.py
MOD = 998244353
def precompute_factorials(n):
# fac: 階乗を格納するリスト。0!と1!は1で初期化
fac = [1] * (n+1)
# inv: 逆元を格納するリスト。0の逆元と1の逆元は1で初期化
inv = [1] * (n+1)
# 階乗の逆元を格納するリスト
finv = [1] * (n+1)
for i in range(2, n+1):
# i!を計算し、MODで割った余りを保存
fac[i] = fac[i-1] * i % MOD
# フェルマーの小定理を使用
inv[i] = pow(i, MOD-2, MOD)
finv[i] = finv[i-1] * inv[i] % MOD # finv[i] = pow(fac[i], MOD-2, MOD) と同義だけど、効率の良さからこちらを選択
return fac, inv, finv
if __name__=="__main__":
n = 10
fac, inv, finv = precompute_factorials(n)
print(f"fac[5] = {fac[5]}")
print(f"fac[5] * finv[5] % MOD = {(fac[5] * finv[5]) % MOD}")
AtCoder360 Dでトラウマになった逆元の計算方法を復習。
定期的に扱う問題が出てくれると苦手意識を払拭できてありがたい。
要素の出現頻度に基づく組み合わせ数の計算
sample.py
from collections import defaultdict
MOD = 998244353
# 2～10までの要素は998244353を法にした逆数
# 748683265*4≡1(mod 998244353)
inv = [1, 1, 499122177, 332748118, 748683265, 598946612, 166374059, 855638017, 873463809, 443664157, 299473306]
def calculate_combinations():
# 入力リスト
a = [1, 1, 2, 2, 2, 3, 3, 3, 3]
# 各要素の出現回数をカウントする辞書
freq = defaultdict(int) # {1: 2, 2: 3, 3: 4}
for num in a:
freq[num] += 1
# 入力リストの長さを取得
max_length = len(a)
# 結果を格納するリスト（初期値は全て0）
ans = [0] * max_length
# 各ユニークな要素の出現回数について処理
for count in freq.values():
x = 1 # 組み合わせの数を計算するための変数
# j+1個選ぶ組み合わせを計算 (0 ≤ j < count)
for j in range(count):
# count個からj+1個選ぶ組み合わせの数を計算
# from itertools import combinations
# len(list(combinations(range(count), j+1)))と同義(剰余を取るため直接実装)
x = x * (count - j) % MOD
x = x * inv[j+1] % MOD
# 計算結果をans[j]に加算（MODで割った余りを取る）
ans[j] = (ans[j] + x) % MOD
input(ans)
return ans[:len(freq.values())+1]
if __name__=="__main__":
result = calculate_combinations()
print(result)
# result[0] = C(2,1)+C(3,1)+C(4,1) = 9
# result[1] = C(2,2)+C(3,2)+C(4,2) = 10
# result[2] = C(3,3)+C(4,3) = 5
# result[3] = C(4,4) = 1
等差数列的な部分列のカウント
sample.py
from collections import defaultdict
MOD = 998244353
def count_sequences():
a = [2, 4, 8, 14, 12, 16]
n = len(a)
# 結果を格納する配列を初期化
ans = [0] * n
# 長さ1の部分列の数（各要素自身）
ans[0] = n
# 長さ2の部分列の数（2点の組み合わせ）
ans[1] = n * (n-1) // 2
# 動的計画法の中間結果を格納する辞書
d = defaultdict(int)
# 最初の点を選ぶループ
for i in range(n-1):
# 2番目の点を選ぶループ
for j in range(i+1, n):
# 同じ値の場合はスキップ
if a[i] == a[j]:
continue
# 2点間の差分を計算
diff = a[j] - a[i]
# 2点で構成される等差数列のキー
d_key = (i, j, 1)
# 2点で構成される等差数列をカウント
d[d_key] = 1
for k in range(j+1, n):
# 等差数列の条件を満たさない場合はスキップ
if (a[k] - a[i]) % diff != 0:
continue
# 等差数列における位置を計算
l = (a[k] - a[i]) // diff
# lが2以上n-1以下であることを確認
## l <= 1 のチェックで、既にカウント済みの2点の組み合わせを除外する。
## l >= n のチェックで、配列の範囲外アクセスを防ぐ。
if l <= 1 or l >= n:
continue
# 新しい状態のキー
new_key = (i, j, l)
# 前の状態のキー
prev_key = (i, j, l-1)
# 前の状態が存在する場合
if prev_key in d:
# 新しい状態に前の状態の数を加算
d[new_key] += d[prev_key]
d[new_key] %= MOD
# 結果配列を更新
ans[l] += d[prev_key]
ans[l] %= MOD
return ans
if __name__=="__main__":
print(count_sequences())
# プログラム想定
# a = [2, 4, 8, 14]
## Step_01
## 2, 4 ... diff:2, d_key: (0,1,1), d[(0,1,1)]=1
## 2, 8 ... diff:6, d_key: (0,2,1), d[(0,2,1)]=1
## 2, 14 ... diff:12, d_key: (0,3,1), d[(0,3,1)]=1
## 4, 8 ... diff:4, d_key: (1,2,1), d[(1,2,1)]=1
## 4, 14 ... diff:10, d_key: (1,3,1), d[(1,3,1)]=1
## 8, 14 ... diff:2, d_key: (2,3,1), d[(2,3,1)]=1
## Step_02(diff:6, d_key: (0,2,1), d[(0,2,1)]=1の場合)
## k: 3 - 3
## 14-2 % 6 = 0なので条件を満たす
## l = 12 // 6 = 2なので、条件を満たす(2項目)
# new_key = (1, 2, 2)
# prev_key = (1, 2, 1)
# d[(1, 2, 2)] += d[(1, 2, 1)]
# ans[2] += d[(1, 2, 1)]
雑記.1
部分列が等差数列になっているか否かの判定がこの問題の鍵なんだけど、少し難しいな。
等差数列部分列判定問題
sample.py
def count_arithmetic_subsequences():
# 入力配列（ここではハードコーディングされている）
A = [2, 5, 7, 9, 12, 15]
N = len(A) # 配列の長さ
count = 0 # 等差数列の総数をカウントする変数
# すべての可能な部分列の開始点と2番目の要素を選ぶ
ans_list = []
for i in range(N):
for j in range(i+1, N):
current_diff = A[j] - A[i] # 現在の等差
subsequence_length = 2 # 部分列の長さ（最初の2要素）
count += 1 # 長さ2の部分列をカウント
ans_list.append([A[i],A[j]]); idx=len(ans_list)-1
# 3番目以降の要素を探索
prev = j # 前の要素のインデックス
for k in range(j+1, N):
# 等差が一致する場合
if A[k] - A[prev] == current_diff:
# 部分列の長さを増加
subsequence_length += 1
# 新しい部分列をカウント
count += 1
# 前の要素を更新
prev = k
ans_list[idx].append(A[k])
else:
# 等差が一致しない場合、この開始点での探索を終了
continue
# 結果の出力
print(f"見つかった等差数列の総数: {count}")
print(f"等差数列: {ans_list}")
if __name__=="__main__":
count_arithmetic_subsequences()
# 見つかった等差数列の総数: 18
# 等差数列: [[2, 5], [2, 7, 12], [2, 9], [2, 12], [2, 15], [5, 7, 9], [5, 9], [5, 12], [5, 15], [7, 9], [7, 12], [7, 15], [9, 12, 15], [9, 15], [12, 15]]
雑記.2
このコードだと、100%TLEするだろうけど、部分列の中の等差数列の求め方はこんな感じかな。
数列内の等差部分列の総数を動的計画法でカウント
sample.py
def count_arithmetic_subsequences(nums):
n = len(nums)
# # インデックス i を終点とし、公差 diff を持つ等差部分列の個数を記録する辞書
dp = [{} for _ in range(n)] # dp[i][diff]
total = 0
for i in range(n):
for j in range(i):
diff = nums[i] - nums[j]
# 要素 i と要素 j の差 diff が、インデックス j を終点とする等差部分列の中に存在していれば、その個数を prev に格納する。
# 存在しなければ、prev に 0 を格納する。
prev = dp[j].get(diff, 0)
# i 番目の要素を見たときに、等差が diff である等差部分列の個数は、
# 現在 i 番目で保持している diff の個数（すでに存在するもの）と、
# j 番目の要素が持っている diff の個数（過去から引き継ぐもの）に、新たに作られる1つのペア (nums[j], nums[i]) を加えたもの。
# -- dp[i].get(diff, 0)
# --- インデックス i を終点とし、公差が diff の等差部分列が現在いくつ存在するかを取得します。
# -- diff
# --- インデックス j を終点とし、公差が diff の等差部分列がいくつ存在するかを取得
# -- +1
# --- 新しく作られるペア (nums[j], nums[i]) 自体も1つの新しい等差部分列としてカウント
dp[i][diff] = dp[i].get(diff, 0) + prev + 1
total += prev + 1
return total, dp
# dp[i][diff] とは？
# 例えば、数列 [2, 4, 6, 8] において、dp[3][2] = 3 であれば、インデックス 3（値が8）を終点とし、公差が2の等差部分列が3つ存在することを意味する。
if __name__=="__main__":
total, dp = count_arithmetic_subsequences([2, 4, 6, 8])
print(dp)
print(total)
雑記.3
動的計画法一時期、集中してやってたけど、完全に身につかなかったなぁ。
このコード理解するだけでも1時間くらいかかった⋯
公差と長さの条件を満たす等差部分数列の個数
sample.py
from collections import defaultdict
def optimized_dp(nums, max_diff=100, threshold=3):
n = len(nums)
dp = [defaultdict(int) for _ in range(n)]
total = 0
# 「j番目の要素」と「i番目の要素」のペアを比較し、それらを基に等差数列を構築
# (外側のループ) i 番目の要素を終点とする等差数列の情報がすべて更新される
for i in range(n):
for j in range(i):
diff = nums[i] - nums[j]
# 差分が max_diff を超える場合はスキップ
if abs(diff) > max_diff:
continue
# j番目の要素で終わり、公差が diff の等差数列の個数
prev = dp[j].get(diff, 0)
# 新たに i 番目の要素を加えることで、長さが1増加した等差数列を構築できるため、+1 をする。
current = prev + 1
# i番目の要素で終わり、公差が diff の等差数列の個数
dp[i][diff] += current
if current >= threshold:
total += 1
# if dp[i][diff] < threshold:
# del dp[i][diff]
# print(dp)
return total
if __name__=="__main__":
# テスト用データ
nums = [3, 1, 4, 1, 5, 9, 2, 6, 5, 3]
max_diff = 10
threshold = 3
# 関数呼び出し
result = optimized_dp(nums, max_diff=max_diff, threshold=threshold)
# 結果表示
print(f"条件を満たす等差数列の総数: {result}")
# 3 {}
# 1 {-2: 1} -> 1を終端とする、交差が-2の部分数列の長さは1
# 4 {1: 1, 3: 1} -> 4を終端とする、公差が1の部分数列の長さは1、公差が3の部分数列の長さは1。
# 1 {-2: 1, 0: 1, -3: 1} -> 1を終端とする、公差が-2の部分数列の長さは1、公差が0の部分数列の長さは1、公差が-3の部分数列の長さは1。
# 5 {2: 1, 4: 2, 1: 2} -> 5を終端とする、公差が2の部分数列の長さは1、公差が4の部分数列の長さは2、公差が1の部分数列の長さは2
# 9 {6: 1, 8: 2, 5: 1, 4: 3} -> 9を終端とする、公差が6の部分数列の長さは1、公差が8の部分数列の長さは2、公差が5の部分数列の長さは1、公差が4の部分数列の長さは3
# 2 {-1: 1, 1: 2, -2: 1, -3: 1, -7: 1} -> 2を終端とする、公差が-1の部分数列の長さは1、公差が1の部分数列の長さは2、公差が-2の部分数列の長さは1、公差が-3の部分数列の長さは1、公差が-7の部分数列の長さは1
# 6 {3: 1, 5: 2, 2: 1, 1: 3, -3: 1, 4: 1} -> 6を終端とする、公差が3の部分数列の長さは1、公差が5の部分数列の長さは2、公差が2の部分数列の長さは1、公差が1の部分数列の長さは3、公差が-3の部分数列の長さは1、公差が4の部分数列の長さは1
# 5 {2: 1, 4: 2, 1: 2, 0: 1, -4: 1, 3: 1, -1: 1} -> 5を終端とする、公差が2の部分数列の長さは1、公差が4の部分数列の長さは2、公差が1の部分数列の長さは2、公差が0の部分数列の長さは1、公差が-4の部分数列の長さは1、公差が3の部分数列の長さは1、公差が-1の部分数列の長さは1
# 3 {0: 1, 2: 2, -1: 1, -2: 2, -6: 1, 1: 3, -3: 2} -> 3を終端とする、公差が0の部分数列の長さは1、公差が2の部分数列の長さは2、公差が-1の部分数列の長さは1、公差が-2の部分数列の長さは2、公差が-6の部分数列の長さは1、公差が1の部分数列の長さは3、公差が-3の部分数列の長さは2