Official

C - Colorful Candies Editorial by en_translator


Let us denote the \(i\)-th, \((i+1)\)-th, \((i+2)\)-th, \(\ldots\), and \((i+K-1)\)-th candies by the segment \(\lbrack i, i+K-1 \rbrack\).

There are \(N+K-1\) candidates of the segment of Takahashi’s choice, \(\lbrack 1, K \rbrack, \lbrack 2, K+1 \rbrack, \lbrack 3, K+2 \rbrack, \ldots,\) and \(\lbrack N-K+1, N \rbrack\).
It is sufficient to find for each segment “the number of colors of candies in it” and output the maximum value as the answer.

In order to count the number of colors of candies in a segment we can “inspect each of the \(K\) candies in the segment and count the colors within it.”
However, this method requires \(K\) times of “inspection of a candy” for each segment, so if we do this for every of \(N - K + 1\) candidates of segments, we need \(K(N - K + 1)\) number of operations of “inspecting a candy.” This method costs to much computational time, exceeding the time limit.

So, we will speed up the method above.
Assume that we have just inspected \(\lbrack i, i+K-1 \rbrack\) and now want to inspect \(\lbrack i+1, i+K \rbrack\) right after that. Here, note that adding the \((i+K)\)-th candy to, and then instead removing the \(i\)-th candy from \(\lbrack i, i+K-1 \rbrack\), results in \(\lbrack i+1, i+K \rbrack\). That is, \(\lbrack i, i+K-1 \rbrack\) and \(\lbrack i+1, i+K \rbrack\) have “little difference.” So, when inspecting \(\lbrack i+1, i+K \rbrack\), instead of inspecting all the candies in \(\lbrack i+1, i+K \rbrack\), we can only focus on the “difference” between the \(\lbrack i, i+K-1 \rbrack\) that just have been checked.

We consider the following algorithm based on the idea.

  1. First we inspect \(c_1, c_2, \ldots,\) and \(c_K\), and for each color of which at least one of candy in \(\lbrack 1, K \rbrack\) has, memorize the number of candies of that color with an associated array (like std::map in C++ language). Here, note that the number of entries in the associated array is equal to the number of distinct colors in \(\lbrack 1, K \rbrack\).
  2. Next, we increment the number for Color \(c_{K+1}\) in the associated array, and decrement the number for Color \(c_1\). (If the number for Color \(c_1\) became \(0\) as a result, we remove the entry for the Color \(c_1\) from the associated array.) Then, the number of entries in the associated array is equal to the number of distinct colors in \(\lbrack 2, K+1 \rbrack\).
  3. Then, for the associated array, we increment the number for Color \(c_{K+2}\), and decrement the number for Color \(c_2\). (If the number for Color \(c_2\) became \(0\) as a result, we remove the entry for the Color \(c_2\) from the associated array.) Then, the number of entries in the associated array is equal to the number of distinct colors in \(\lbrack 3, K+2 \rbrack\).
  4. By repeating this, for each of \(\lbrack 1, K \rbrack, \lbrack 2, K+1 \rbrack, \lbrack 3, K+2 \rbrack, \ldots, \lbrack N-K+1, N \rbrack\), we can obtain the number of distinct colors in the segment as the number of entries in the associated array, so all that left is output the maximum value of them.

For each candy, the addition to and removal from the associated array is done at most \(N\) times, so the time complexity for this algorithm is \(\mathrm{O}(N\log N)\), which is fast enough to get accepted for this problem.

The sample code in C++ language for this problem follows.

#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: