【競技プログラミング】AtCoder Beginner Contest 146_E問題

問題

要約

1. 長さNの正整数列A1, A2, ..., ANと正の整数Kが与えられる。

2. Aの空でない連続する部分列について、以下の条件を満たすものの数を求める。

   • 部分列の要素の和をKで割った余りが、その部分列の要素の数と等しくなる

3. 2つの部分列が同じ要素を持っていても、元の配列内での位置が異なれば別のものとして数える。

変数情報：

• N: 正整数列の長さ
• A1, A2, ..., AN: 与えられる正整数列
• K: 正の整数（除数として使用）

既存投稿一覧ページへのリンク

一覧ページ

独り言

正直、こういう数学系の問題を中学生や高校生が平気で解けるというのが凄いよなぁ。
まぁ、部屋掃除している時に、受験時代のノートを開いたりすると、わけのわからん数式を大量に書いてあるのを見て、当時の私は本当に理解していたのか？と驚くこともあるんだけど・・・
何を言いたいかと言うと、学生ってなんだかんだ言って優秀だとAtCoderを解くたびに思いしらされますよ。

アプローチ

累積和と余りの性質を利用して、効率的に条件を満たす部分列の数をカウントする。
配列の各要素を前処理し、累積和の余りを計算することで、任意の部分列の和の余りを高速に求められるようにする。
その後、ハッシュマップを使用して条件を満たす部分列の数をカウントする。

解法手順

1. 入力を受け取り、配列Aを初期化する。

2. 配列Aの各要素を前処理する：

   • A[0]を(A[0]-1) % Kに変更する。
   • A[i]を(A[i] + A[i-1] - 1) % Kに変更する（i > 0の場合）。

3. 答えを格納する変数ansを0で初期化する。

4. ハッシュマップbを作成し、0をキーとして1を値として初期化する。

5. 配列Aを順に走査し、以下の処理を行う：

   • i > 0の場合、b[A[i-1]]の値を1増やす（存在しない場合は新たに作成）。
   • i >= Kの場合、b[A[i-K]]の値を1減らす。
   • i == K-1の場合、b[0]の値を1減らす。
   • b[A[i]]の値（存在しない場合は0）をansに加算する。

6. 最終的なansの値を出力する。

ACコード

ac.py

def io_func():
    # 入力を受け取る
    n, k = map(int, input().split())  # nとkを入力から取得
    a = list(map(int, input().split()))  # 配列Aを入力から取得
    return n, k, a

def solve(n, k, a):
    # 配列Aの前処理
    a[0] = (a[0] - 1) % k  # 最初の要素を処理
    for i in range(1, n):
        a[i] = (a[i] + a[i-1] - 1) % k  # 累積和の余りを計算

    ans = 0  # 答えを格納する変数を初期化
    b = {0: 1}  # ハッシュマップbを初期化

    for i in range(n):
        if i > 0:
            # i-1番目の要素の出現回数を増やす
            b[a[i-1]] = b.get(a[i-1], 0) + 1
        if i >= k:
            # i-k番目の要素の出現回数を減らす
            b[a[i-k]] -= 1
        if i == k - 1:
            # k-1番目の要素の時、0の出現回数を減らす
            b[0] -= 1
        # 現在の要素の出現回数をansに加算
        ans += b.get(a[i], 0)

    return ans  # 最終的な答えを返す

if __name__=="__main__":
    # メイン処理
    n, k, a = io_func()  # 入力を受け取る
    result = solve(n, k, a)  # 解を計算
    print(result)  # 結果を出力

###
# n: 配列Aの要素数
# k: 条件となる数
# a: 入力配列A（前処理後は累積和の余りを格納）
# ans: 条件を満たす部分列の数
# b: 各余りの出現回数を格納するハッシュマップ

# 1. io_func関数: 標準入力から値を読み取り、n, k, aを返す
# 2. solve関数: 主要なロジックを実装
#    - 配列aの前処理: 累積和の余りを計算
#    - ハッシュマップbを使用して、各余りの出現回数を管理
#    - 条件を満たす部分列の数をカウント
# 3. メイン処理: io_func関数でデータを取得し、solve関数で解を計算、結果を出力
# 
# このプログラムは、累積和と余りの性質を利用して効率的に条件を満たす部分列の数をカウントしています。
# ハッシュマップを使用することで、O(n)の時間複雑度で解を求めることができます。```

思考過程

1. 累積和の利用：
   部分列の和を高速に計算するために累積和を使用する。
2. 余りの性質：

(A+B)modK=((AmodK)+(BmodK))modK

という性質を利用して、累積和の余りを効率的に計算します。
3. ハッシュマップの活用
各余りの出現回数を管理するためにハッシュマップを使用し、O(1)での検索と更新を可能にする。
4. スライディングウィンドウ
長さKの範囲内で条件を満たす部分列を効率的にカウントするために、スライディングウィンドウの考え方を適用する。

計算量の概算

• 配列の前処理：O(n)
• メインのループ処理：O(n)

その他周辺知識

整数の除法と剰余（モジュロ演算）

a[0] = (a[0] - 1) % k
a[i] = (a[i] + a[i-1] - 1) % k

ある数を別の数で割った余りを求める演算。
数学的には、

a \equiv b \pmod{m}

と表現される。
この操作により、全ての要素を0からk-1の範囲に収める。
リストの最初の要素（インデックス0）を更新。

累積和

for i in range(1, n):
    a[i] = (a[i] + a[i-1] - 1) % k

この部分で累積和を計算している。

ハッシュマップを用いた要素の出現回数の管理

b = {0: 1}
...
b[a[i-1]] = b.get(a[i-1], 0) + 1
...
b[a[i-k]] -= 1
...
ans += b.get(a[i], 0)

ハッシュマップbを使用して、各余りの出現回数を管理する。
辞書bを初期化し、キーが存在しない場合のデフォルト値を0として、値を1増やしている。

数え上げ

ans = 0
...
ans += b.get(a[i], 0)

区間と差分

if i >= k:
    b[a[i-k]] -= 1

長さkの区間を考え、その区間の開始点（i-k）が範囲外に出たときに、その要素の出現回数を減らす。