AtCoderログ：0024 - ABC 210 C

問題：ABC 210 C - Colorful Candies

問題文

$N$ 個のキャンディが左右一列に並んでいます。
それぞれのキャンディは、色 $1$、色 $2$、$\ldots$、色 $10^9$ の、$10^9$ 種類の色のうちいずれかの色をしています。
$i=1,2,\ldots,N$ について、左から $i$ 番目のキャンディの色は色 $c_i$ です。
高橋君は並んでいるキャンディのうち、連続して並んだ $K$ 個のキャンディをもらうことができます。
すなわち、$1 \le i \le N−K+1$ を満たす整数 $i$ を選んで、 左から $i$ 番目、$i+1$ 番目、$i+2$ 番目、$\ldots$ 、$i+K−1$ 番目のキャンディをもらうことができます。
高橋君はいろいろな色のキャンディを食べたいので、もらうキャンディに含まれる色の種類数が多いほどうれしい気持ちになります。
? 高橋君がもらうキャンディに含まれる色の種類数の最大値を出力してください。

制約

・$1 \le K \le N \le 3 \times 10^5$
・$1 \le c_i \le 10^9$
・入力はすべて整数

回答1 (TLE)

キャンディの色を配列 color() に保存することにします。処理の内容は、各 i に対して左から i 番目から i+k−1 番目の k 個の配列値から色の種類を調べ、その最大値を求めるという流れになります。色の種類を計算するには、仮想配列 candy を使用しましょう。candy は key に整数値、value にブール値をとることにして、j=i, i+1, ..., i+k-1 に対して candy[color.at(j)]=true とおいていけば、色の種類は candy の要素数 (int)candy.size() として求められます。

コードは以下のようになりましたが、TLEとなってしまいました。このコードの計算量は $O((N-K+1)K)$ であるため、時間オーバーになってしまったのでしょう。

abc210c-1.cpp

#include <bits/stdc++.h>
using namespace std;
 
int main() {
  int n, k;
  cin >> n >> k;
 
  vector<int> color(n);
  for ( int i=0; i<n; i++ ) {
    cin >> color.at(i);
  }
 
  int cmax = -1;
  for ( int i=0; i<n-k+1; i++ ) {
    map<int,bool> candy;
    for ( int j=i; j<i+k; j++ ) {
      candy[color.at(j)] = true;
    }
    cmax = max( cmax, (int)candy.size() );
  }
  cout << cmax <<endl;
}

回答2 (AC)

回答1の無駄な処理を削減しましょう。回答1で最も無駄なのは、j=i, i+1, ..., i+k-1 に対して仮想配列を構成して要素数を求める処理と、その次の j=i+1, i+2, ..., i+k-1, i+k に対して仮想配列を構成して要素数を求める処理が独立に行われていることです。そこでこの2つの処理を関連付けます。

今、j=i, i+1, ..., i+k-1 に対する仮想配列 candy が求まっているとき、ここから j=i+1, i+2, ..., i+k-1, i+k に対する仮想配列を求めることを考えます。
candy は color.at(i), color.at(i+1), ..., color.at(i+k-1) の重複をなくした値の集合なので、処理としては (1) color.at(i) に関する情報を削除し、(2) color.at(i+k) に関する情報を追加する、の2つになります。

(1)は candy から color.at(i) を削除すれば良さそうですが、それではうまくいきません。他の color.at(i') で同じ色を使っている可能性があるからです。そこで仮想配列 candy の仕様を、key は色を表す int のまま、value は何回登場したかを表すカウンタを表す int に変更します。こうすることで、candy(color.at(i)) の値を確認することで、削除すべきかを判定します。具体的には、candy(color.at(i))>1 ならばカウンタ値を 1 減らし、candy(color.at(i))==1 ならその key を削除します。
(2)は candy に color.at(i+k) を追加することになります。具体的には、candy(color.at(i+k)) が既に存在していればその値を 1 増やし、なければその key を追加します。なお、上記により candy の要素が増減する場合が把握出来たので、candy に含まれる色の種類を表す変数 ncolor を用意し、その値を管理するようにしています。これにより、candy.size を実行する必要がなくなりました。

以上の考察を反映させたコードは以下のようになりました。計算量は $O(K+N-K)=O(N)$ となるので、無事、ACとなりました。

abc210c-2.cpp

#include <bits/stdc++.h>
using namespace std;
 
int main() {
  int n, k;
  cin >> n >> k;
 
  vector<int> color(n);
  for ( int i=0; i<n; i++ ) {
    cin >> color.at(i);
  }
 
  map<int,int> candy;
  for ( int i=0; i<k; i++ ) {
    if ( candy.count(color.at(i)) ) {
      candy[color.at(i)] += 1;
    } else {
      candy[color.at(i)] =  1;
    }
  }
 
  int ncolor = int(candy.size()), cmax = ncolor;
  for ( int i=0; i<n-k; i++ ) {
    if ( candy.count(color.at(i)) ) {
      if ( candy[color.at(i)]>1 ) {
        candy[color.at(i)] -= 1;
      } else {
        candy.erase(color.at(i));
        ncolor -= 1;
      }
    }
 
    if ( candy.count(color.at(i+k)) ) {
      candy[color.at(i+k)] += 1;
    } else {
      candy[color.at(i+k)] =  1;
      ncolor += 1;
    }
    cmax = max( cmax, ncolor );
  }
 
  cout << cmax <<endl;
}

調べたこと

AtCoder の解説 → 公式解説

回答2と同じ方式ですが、candy の色の種類は size 関数を使用していました。

AtCoder の解説 → ユーザ解説

回答2と同じ方式でした。回答2のようなアプローチを尺取り法と呼ぶそうです。計算量評価が $O(N \log{N})$ となっていましたが、何を単位に数えているかわかりませんでした。

リンク

• 前の記事 → AtCoderログ：0021 - ABC 210 B
• 次の記事 → AtCoderログ：0023 - ABC 102 B