長さKの固定窓は、右へ1つ進むたびに右側の色を1つ入れ、左端の色を1つ外して更新します。ABC210 Cのような色の頻度表を保つ処理では、色の正負や大小に基づく単調性は必要ありません。
一方、窓の長さを条件に合わせて伸縮する尺取り法には、右端を進めた後に左端を後ろへ戻さずに進められる単調性が必要です。負数を含む合計条件ではこの性質が崩れるため、同じ手順を一律には使えません。
負数を含む合計条件で単調性が崩れる
合計が5以上となる最短の連続区間を探す例を見ます。数列[1, -100, 5]で左端を固定した合計は、右へ進むにつれて1、-99、-94です。正数列を前提にする典型的な尺取り法では、現在の合計が5未満なら左端を縮めません。そのため右端まで進んでも条件を満たせず、末尾だけの区間[5]（合計5）を見つけられません。負数を加えると合計が下がり得るため、右を進めれば条件を満たすとは限らない例です。
ABC210 C – Colorful Candiesの問題
AtCoder公式問題では、色c_iのキャンディがN個並んでいます。連続するK個を選んだときの色の種類数について、取り得る区間の最大値を求めます。開始位置は1始まりで1 ≤ i ≤ N-K+1です。
標準入力の1行目はN K、2行目はc_1 c_2 ... c_Nです。出力するのは、長さKの連続区間に含まれる異なる色の最大数です。
1 ≤ K ≤ N ≤ 3 * 10^51 ≤ c_i ≤ 10^9- 入力値はすべて整数です。
公式サンプルと標準入力での実行結果
次の公式サンプル3件を掲載コードへそのまま標準入力し、公式出力とローカル実行結果を照合しました。ローカル実行はAtCoderへの提出やACを意味しません。
| 公式サンプル | 標準入力 | 公式出力 | 掲載コードの実行結果 |
|---|---|---|---|
| 1 |
|
3 |
3 |
| 2 |
|
1 |
1 |
| 3 |
|
4 |
4 |
右を入れて左を外し、頻度表を更新する
AtCoder公式解説も、最初のK個を数えた後、右から入る色を加え、左から出る色を減らして窓を更新します。色の頻度が0になったときは、その色のkeyを削除します。現在の窓だけをfreqに正確に残せば、keyの数がその窓の異なる色数です。
手計算例としてN=5、K=3、色列[1, 2, 1, 3, 3]を追います。追加直後は一時的にK+1個ですが、左端を外した後に長さKの窓ができ、その時点で答えを更新します。
| 段階 | 窓または一時状態 | 操作 | 更新後のfreq | 種類数 | answer |
|---|---|---|---|---|---|
| 初期窓 | [1, 2, 1] |
最初の3個を数える | {1: 2, 2: 1} |
2 | 2 |
| 1回目・追加直後 | [1, 2, 1, 3] |
右の3を加える |
{1: 2, 2: 1, 3: 1} |
3（まだ判定しない） | 2を保持 |
| 1回目・削除後 | [2, 1, 3] |
左の1を1減らす |
{1: 1, 2: 1, 3: 1} |
3 | max(2, 3)=3 |
| 2回目・追加直後 | [2, 1, 3, 3] |
右の3を加える |
{1: 1, 2: 1, 3: 2} |
3（まだ判定しない） | 3を保持 |
| 2回目・削除後 | [1, 3, 3] |
左の2が0になり、keyを削除 |
{1: 1, 3: 2} |
2 | max(3, 2)=3 |
最後の窓では色2の頻度が0になったため、freqからkeyごと削除しています。初期窓では種類数2、その後の窓では3と2を得るので、すべての窓で最大の種類数は3です。
Python 3の完全コード
公式問題の色列は1始まりで説明されますが、Pythonのリストcolorsは0始まりです。rightは右から入る要素の添字で、追加後にcolors[right-k]を外すと、更新後の窓はcolors[right-k+1:right+1]になります。
import sys
def solve():
data = list(map(int, sys.stdin.buffer.read().split()))
n, k = data[0], data[1]
colors = data[2:2 + n]
freq = {}
for color in colors[:k]:
freq[color] = freq.get(color, 0) + 1
answer = len(freq)
for right in range(k, n):
entering = colors[right]
freq[entering] = freq.get(entering, 0) + 1
leaving = colors[right - k]
freq[leaving] -= 1
if freq[leaving] == 0:
del freq[leaving]
answer = max(answer, len(freq))
print(answer)
if __name__ == '__main__':
solve()
freqの不変条件から正しさを説明する
初期化の後と、左右の更新を終えた後には、freqのkeyは現在の長さKの窓にある色だけで、各valueはその窓での正確な出現回数です。右の色を1つ加え、左の色を1つ減らすことで次の窓へ移り、0になったkeyを削除すれば、この不変条件は保たれます。更新途中の追加直後だけは一時的にK+1個です。
keyに0回の色が残らないので、窓の異なる色数はlen(freq)と一致します。最初の窓でanswerを初期化し、左右を1つずつずらした後の各窓でもlen(freq)を比較するため、すべての窓を一度ずつ評価して最大値を取れます。初期窓の開始位置は0で、rightをKからN-1まで動かすと、残りの開始位置0始まりの1からN-Kまでを順に評価します。
計算量
初期窓のK個を一度ずつ数え、以後はN-K回の移動で各回1回追加・1回削除します。Pythonのdict操作が平均O(1)であることを前提に、時間計算量は期待・平均O(N)です。頻度表に入るkeyは窓の異なる色だけなので、freqの領域はO(K)です。標準入力を保持するdataと色列colorsも持つため、入力配列を含むプログラム全体のメモリはO(N)です。
AtCoder公式解説の掲載例はC++のstd::mapで、解説では時間計算量をO(N log N)としています。ここでのPythonの期待・平均O(N)は、dictの平均操作時間を前提にした別実装の見積もりです。
区間問題の対照として、ABC143 Dはソート後の二分探索であり、二ポインタ実装ではありません。ここでは方法の違いだけを示し、別問題のコードは載せません。
独自例で境界を確認する
次の3件は公式サンプルではなく、K=1、K=N、重複色の頻度と0件keyの削除を確認する独自例です。掲載コードを各標準入力で実行したローカル結果を示します。これもAtCoderへの提出やACを意味しません。
| 確認するケース | 標準入力 | 掲載コードの実行結果 |
|---|---|---|
K=1 |
|
1 |
K=N |
|
2 |
| 重複色と0件key削除 |
|
3 |
固定長窓では、各移動で入る要素と出る要素を対応させ、必要な状態を正確に更新できるかを考えます。窓の長さ自体を条件に合わせて変える問題では、左端を戻さず進められる性質があるかを別に確かめてください。
