Official

C - Colorful Candies Editorial by leaf1415


\(i\) 番目、\(i+1\) 番目、\(i+2\) 番目、\(\ldots\)\(i+K-1\) 番目のキャンディをまとめて、区間 \(\lbrack i, i+K-1 \rbrack\) で表すことにします。

高橋君がとる区間の候補は、 \(\lbrack 1, K \rbrack, \lbrack 2, K+1 \rbrack, \lbrack 3, K+2 \rbrack, \ldots, \lbrack N-K+1, N \rbrack\)\(N+K-1\) 個があります。
これらすべてについて「区間に含まれるキャンディの色の種類数」を求め、それらの最大値を答えとして出力すれば良いです。

\(1\) つの区間に含まれるキャンディの色の種類数を求めるには 「区間内の \(K\) 個のキャンディを一つずつ見て、区間内の色の種類数を数える」という方法が考えられます。
しかしこの方法では、\(1\) つの区間を調べるのに「キャンディを見る」操作を \(K\) 回行うので、\(N-K+1\) 個の候補区間すべてにこれを行うと、 アルゴリズム全体で「キャンディを見る」操作を \(K(N-K+1)\) 回行うことになります。この方法では時間計算量が大きすぎるため、実行時間制限に間に合いません。

そこで、上記の方法を高速化します。
いま、\(\lbrack i, i+K-1 \rbrack\) を調べてその直後に \(\lbrack i+1, i+K \rbrack\) を調べたいとします。 ここで、\(\lbrack i, i+K-1 \rbrack\) から \(i+K\) 番目のキャンディを追加して代わりに \(i\) 番目のキャンディを削除すると、\(\lbrack i+1, i+K \rbrack\) になることに注意します。 つまり、\(\lbrack i, i+K-1 \rbrack\)\(\lbrack i+1, i+K \rbrack\) には「わずかな差」しかありません。 そこで \(\lbrack i+1, i+K \rbrack\) について調べるときは \(\lbrack i+1, i+K \rbrack\) のキャンディをすべて見るのではなく、 直前に調べた \(\lbrack i, i+K-1 \rbrack\) との「差分」のみに着目することで、高速化を図ることができます。

この考え方に基づいて、以下のアルゴリズムを考えます。

  1. まず \(c_1, c_2, \ldots, c_K\) を見て、 \(\lbrack 1, K \rbrack\) 内にキャンディが \(1\) 個以上存在する各色についてキャンディの個数を連想配列(C++言語ではstd::mapなど)で記録します。このとき、連想配列のエントリ数は \(\lbrack 1, K \rbrack\) 内の色の種類数に等しいことに注意します。
  2. 次に、連想配列内の色 \(c_{K+1}\) の個数を \(1\) 増やし、色 \(c_1\) の個数を \(1\) 減らす操作を行います(その結果、色 \(c_1\) の個数が \(0\) 個になったら連想配列から色 \(c_1\) のエントリを削除します。)すると、連想配列のエントリ数は \(\lbrack 2, K+1 \rbrack\) 内の色の種類数に等しくなります。
  3. さらに、連想配列に対して、色 \(c_{K+2}\) の個数を \(1\) 増やし、色 \(c_2\) の個数を \(1\) 減らす操作を行います。(その結果、色 \(c_2\) の個数が \(0\) 個になったら連想配列から色 \(c_2\) のエントリを削除します。)すると、連想配列のエントリ数は \(\lbrack 3, K+2 \rbrack\) 内の色の種類数に等しくなります。
  4. これを繰り返すことで、\(\lbrack 1, K \rbrack, \lbrack 2, K+1 \rbrack, \lbrack 3, K+2 \rbrack, \ldots, \lbrack N-K+1, N \rbrack\) のそれぞれについて、区間内の色の種類数を連想配列のエントリ数として求めることができるので、それらの最大値を答えとして出力すれば良いです。

各キャンディについて、連想配列への追加と削除はそれぞれ高々 \(1\) 回ずつしか行われないので、このアルゴリズムの時間計算量は \(\mathrm{O}(N\log N)\) であり、この問題に正解するためには十分高速です。

以下に、C++言語によるこの問題の正解例を記載します。

#include <iostream>
#include <map>

using namespace std;

int n, k;
int c[300005];
map<int, int> mp;

int main(void)
{
  cin >> n >> k;
  for(int i = 1; i <= n; i++) cin >> c[i];

  for(int i = 1; i <= k; i++) mp[c[i]]++;
  int ans = mp.size();

  for(int i = k+1; i <= n; i++){
    mp[c[i]]++;
    mp[c[i-k]]--;
    if(mp[c[i-k]] == 0) mp.erase(c[i-k]);
    ans = max(ans, (int)mp.size());
  }
  cout << ans << endl;

  return 0;
}

posted:
last update: