はまやんはまやんはまやん

hamayanhamayan's blog

Colorful Candies [AtCoder Beginner Contest 210 C]

https://atcoder.jp/contests/abc210/tasks/abc210_c

前提知識

解説

https://atcoder.jp/contests/abc210/submissions/24333038

色々解法があるが、特別なデータ構造が必要ないのは尺取り法(正確にはっぽいアルゴリズム)だろう。
この問題は尺取り法の入門問題として適切であるので、アルゴリズムも含めて説明しておく。
ただ、有名所で説明が各所でなされているので、図解されている記事を探して読んだ方が手っ取り早いかもしれない。

尺取り法とは

尺取り法と言っているが、基本的には差分計算である。
[1,K]の区間についての計算結果から[2,K+1]の区間についての計算結果を差分だけを計算することで高速に求める
というのがモチベーションとなる。
尺取り法はもうちょっと広い解釈ができるのだが、差分計算だけを行うことで区間をうまく動かしていくという共通認識がある。

ちょっと、ややこしく説明しすぎた気もする。問題に合わせた具体的な説明に戻ろう。

差分計算をする

今回はすべての区間を確認して、種類数の最大値を求める問題であり、実際これをそのまま行っていく。
この際に、すべての区間で種類数をただ求めるには時間がかかるので、差分計算だけを行ってすべての区間を確認していく。

では、最初の区間である[1,K]について計算していこう。
この区間の種類数を計算するために以下を用意する。

kinds := 今見ている区間の種類数
cnt[x] := 今見ている区間に数xが何個含まれているか

これを構築する方法は、[1,K]の数を見て、cnt[x]をインクリメントしていく。
cnt[x]が0からインクリメントされるタイミングだけkindsをインクリメントするようにすれば種類数も数えられる。

さて、この状態から差分計算をして[2,K+1]を計算しよう。
差分としては1番目の要素が取り除かれて、K+1番目の要素が追加されている。
1番目の要素についてcnt[x]を1つ減らして、その結果cnt[x]が0になるなら種類数が1つ減るのでkindsを1つ減らす。
次に、K+1番目の要素についてcnt[x]を1つ増やして、cnt[x]が0からインクリメントされているなら種類数が1つ増えるのでkindsを1つ増やす。
こうすると2回の計算だけで[2,K+1]の場合のkindsとcntを構築することができた。

これをどんどんやっていくと1つ横にずらすための計算はO(1)で済むため、すべての区間を確認してもO(N)で間に合う。
あとは、各区間のkindsについて最大値を取れば答え。

尺取り法とは

一般には区間の右端を1つ動かしたら条件を満たすように左端を右に動かして調整するアルゴリズムのことなのだが、
今回のテクもそれほど変わりないので一緒として説明した。
なので、この問題を理解した後に尺取り法の違う記事や問題も解いてみると、色々なバリエーションが見られるだろう。

int N, K, c[301010];
//---------------------------------------------------------------------------------------------------
void _main() {
    cin >> N >> K;
    rep(i, 0, N) cin >> c[i];

    map<int, int> cnt;
    int kinds = 0;
    int ans = 0;

    rep(i, 0, K) {
        if (cnt[c[i]] == 0) kinds++;
        cnt[c[i]]++;
    }
    chmax(ans, kinds);
    
    rep(i, K, N) {
        if (cnt[c[i - K]] == 1) kinds--;
        cnt[c[i - K]]--;
        if (cnt[c[i]] == 0) kinds++;
        cnt[c[i]]++;
        chmax(ans, kinds);
    }

    cout << ans << endl;
}