ABC 401のC問題 K-bonacci

目次

AtCoder ABC 401 の C 問題を時間内には回答できなかったが、簡単そうなのでチャレンジしてみました。方法は簡単だけど、計算時間制限をクリアできず、そこはカンニング。


AtCoder Beginner Contest 401が 4 月 12 日に開催されました。A と B 問題は簡単に回答できましたが、C 問題は問題分の理解を間違えて正解できませんでした。しかし問題を正しく理解できれば方法自体は簡単そうに思えたので、コンテスト終了後に正解に持っていけました。ただ、計算時間を短縮することが出来ず、そこは既存の回答を参考にさせてもらいました。

問題

C 問題 K-bonacciは、フィボナッチ数列のように前の項目を合計して得られる数列で特定の項の値を求める問題です。今回は、フィボナッチ数列とことなり、考慮する項の数が可変の場合です。

ただし回答では、10^9 で割った余りで答えることとなっています。

方法

フィボナッチ数列では、特定の項の値を求める漸化式が有るようですが、今回の問題に適応できる公式はないようなので正攻法で順番に目的の項まで計算することにしました。ただし前 K 項目の合計を毎回計算するのは避けるという感が働いたいので、過去の計算結果を List 型にメモしておいて、それを使用することとしました。

Ai は、Ai-1 … Ai-k の合計です。しかし毎回この合計を計算し直す必要はありません。Ai-2 より前の合計は計算済みで、その値が Ai-1 になっています。ただし Ai-1 には余分な Ai-k-1 も加算されてしまっているので、その分を引く必要があります。その結果、

$$ Ai = A_{i-1} + A_{i-1} - A_{i-1-k} = 2 \times A_{i-1} - A_{i-1-k} $$

となります。

技術的工夫

List 操作

Ai-k-1 を後で引き算するためには、使用するまでメモしておく必要があります。そこで大きさ K+1 の List 型変数を FIFO の queue として使用しようとしましたが、List 型の操作に時間がかかり制限時間オーバーとなってしまいました。

そこで、メモリー使用量を考えず全ての値を記録できる大きさの List を用意することで List 操作のオーバーヘッドをなくし回答できました。

n, k = map(int, input().split())

if k > n:
    # K could be larger then N.
    print(1)
    exit()

a = [1] * (n + 1)
a[k] = k

for i in range(k + 1, n + 1):
    prev = a[i - 1]
    shift_out = a[i - 1 - k]
    sum = prev * 2 - shift_out
    # print(f"i: {i} -> {sum}")
    a[i] = sum % 1_000_000_000
print(a[n])

全ての値を List に保存するのではなく、必要最小限の List 型を mod した値でアクセスするように修正をしてみましたが、予想に反してメモリー使用量に大きな差は見られませんでした。

n, k = map(int, input().split())

if k > n:
    # K could be larger then N.
    print(1)
    exit()

a = [1] * k
a.append(k)

MOD = len(a)

if k == n:
    print(k)
    exit()

for i in range(k + 1, n + 1):
    prev = a[(i - 1) % MOD]
    shift_out = a[i % MOD]
    sum = prev * 2 - shift_out
    # print(f"i: {i} -> {sum}")
    a[i % MOD] = sum % 1_000_000_000
print(a[i % MOD])

10^9 の余り

回答は、Ai の値をそのままではなく 10^9 で割った余りで回答することになっています。

これは、最後に割り算をするのではなく 10^9 で桁溢れすると考えればよいのだと思いつき、最初から 10^9 以下の部分だけを考えるようにしました。

感想

これまでいくつか C 問題にチャレンジしましたが、どうも自分は問題文を理解する事自体に難があるようです。この問題も時間制限の解決に時間がかかりましたが、問題自体は簡単なのに問題文を間違って解釈してサンプルデータを正解できずに終わってしまいました。

参考サイトと脚注