Official

D - スタンプラリー / Stamp Rally Editorial by admin

gpt-6-astra-medium

概要

連続する区間のうち、全 \(7\) 種類のスタンプを含むものの個数を求めます。区間の右端を順に動かし、各種類が最後に現れた位置から、条件を満たす左端の個数を数えます。

考察

すべての区間を調べると遅い

区間 \((l,r)\) は全部で \(\frac{N(N+1)}{2}\) 個あります。すべての区間を調べる方法では、判定を高速化しても \(O(N^2)\) 時間かかり、\(N \le 10^6\) では間に合いません。

そこで、右端 \(r\) を固定したとき、条件を満たす左端 \(l\) が何個あるかをまとめて求めます。

各種類の「最後に現れた位置」に注目する

チェックポイント \(1\) から \(r\) までを見たとき、種類 \(k\) のスタンプが最後に現れた位置を \(\mathrm{last}[k]\) とします。まだ一度も現れていない種類については \(0\) とします。

区間 \([l,r]\) に種類 \(k\) が含まれるための必要十分条件は、

\(l \le \mathrm{last}[k]\)

です。最後の出現位置が \(l\) 以上ならそのスタンプを区間内で押せますが、\(l\) 未満なら区間内には存在しません。

したがって、全 \(7\) 種類が含まれる条件は、

\(l \le \min_{1 \le k \le 7}\mathrm{last}[k]\)

となります。この最小値を \(m\) とすると、有効な左端は \(1,2,\dots,m\) の \(m\) 個です。

まだ現れていない種類があれば \(m=0\) となるので、有効な区間がない場合も同じ式で扱えます。

具体例

スタンプの並びが次の場合を考えます。

\(1,2,3,4,5,6,7,1\)

右端が \(r=7\) のとき、各種類の最後の出現位置は、

\((1,2,3,4,5,6,7)\)

です。最小値は \(1\) なので、有効な左端は \(l=1\) だけです。

右端を \(r=8\) に進めると、種類 \(1\) の最後の出現位置が更新され、

\((8,2,3,4,5,6,7)\)

になります。最小値は \(2\) なので、今度は \(l=1,2\) の \(2\) 個が有効です。

このように、右端ごとに最小値を答えに加えれば、すべての有効な区間を重複なく数えられます。

アルゴリズム

コードでは種類 \(1,\dots,7\) を添字 \(0,\dots,6\) に対応させます。

  1. 各種類の最後の出現位置を表す配列 last を、すべて \(0\) で初期化します。
  2. last の最小値を表す minimum と、答え answer を \(0\) で初期化します。
  3. 右端 \(i\) を \(1\) から \(N\) まで順に動かします。
    • 現在の種類に対応する添字を c とします。
    • 更新前の last[c] を previous に保存します。
    • last[c] を現在位置 \(i\) に更新します。
    • 必要なら minimum を更新します。
    • answer に minimum を加えます。
  4. answer を出力します。

最小値を再計算するタイミング

毎回 min(last) を求めても、要素数は \(7\) なので十分高速です。提示コードでは、さらに不要な再計算を省いています。

last[c] は、更新によって必ず大きくなります。

  • previous > minimum の場合
    最小値ではない要素を大きくするだけなので、最小値は変わりません。
  • previous == minimum の場合
    最小値を取っていた要素が変わるため、min(last) を再計算します。

これにより、minimum は常に last の最小値に保たれます。各右端について正しい区間数を加算するため、最終的な answer が求める総数になります。

計算量

  • 時間計算量: \(O(7N)=O(N)\)
    各チェックポイントで、最大でも長さ \(7\) の配列の最小値を求めます。
  • 空間計算量: アルゴリズムの補助領域は \(O(7)=O(1)\)
    ただし、提示コードは入力を read().split() で一括保持するため、入力データを含めた使用領域は \(O(N)\) です。

実装のポイント

  • チェックポイントの位置は \(1\) 始まりにします。
    未出現を \(0\) で表せるうえ、最小出現位置がそのまま有効な左端の個数になります。

  • 更新前の値を保存します。
    最小値の再計算が必要か判断するため、last[c] を書き換える前に previous に保存します。

  • 答えの整数型に注意します。
    答えは最大で \(\frac{N(N+1)}{2}\) 以下となり、32 bit 整数型に収まらないことがあります。Python の整数はそのまま扱えますが、C++ などでは 64 bit 整数型が必要です。

    ソースコード

import sys


def main():
    tokens = iter(sys.stdin.buffer.read().split())
    n = int(next(tokens))

    last = [0] * 7
    answer = 0
    minimum = 0

    for i in range(1, n + 1):
        c = int(next(tokens)) - 1
        previous = last[c]
        last[c] = i
        if previous == minimum:
            minimum = min(last)
        answer += minimum

    print(answer)


if __name__ == "__main__":
    main()

この解説は gpt-6-astra-medium によって生成されました。