けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 161 E - Yutori (青色, 500 点)

発想は

とかに似てる。

問題へのリンク

問題概要

長さ  N の o と x で構成された文字列  S が与えられる。 S の index から  K 個選ぶ方法のうち

  • 選んだ index はすべて o である
  • 選んだ index について、どの隣接する 2 つの index についても、その隙間は  C 以上の長さがある (index の差が  C+1 以上である)

という条件を満たすものをすべて考える。そのすべての index の選び方に含まれる index をすべて求めよ。

制約

  •  1 \le K \le N \le 2 \times 10^{5}
  •  0 \le C \le N

考えたこと

これまた問題文がわかりづらいけど、要するに


各 k に対して、S[ k ] = 'o' であるとき、それを 'x' にしたときに「隣接 index の差が  C+1 以上であるような  K 個の 'o' を選ぶ」ことができなくなるかどうかを判定せよ


という問題ということになる。k 番目の o を潰したときに  K 個選ぶことができなくなってしまうならば、その o は絶対に必要だったということがわかる、という論法だ。

まずは、元々の問題で最大で何個の o を選ぶことができるのかを求めておこう。

選べる o の最大個数

これは単純に Greedy で良い。左から順番に o を選んで行って、「最後に選んだ o から C+1 以上離れた o にうち最も左のものを選ぶ」という Greedy で OK。

ここまでで、 O(N^{2}) の解法は得られた。つまり、各 k に対して、S[ k ] を x にしたときの最適解を求めて、それが  K 以上かどうかを判定すればよいのだ。

各 k についての問題を高速に求めたい

各 k について、k 番目のマスを除いた左右の部分について、それぞれ最適解を求めて、それらの合計が  K 未満になれば、k 番目は選ばないといけない、ということがわかる。

よって、

  • left[ i ] := 左から i 文字分についての最適解 (o の個数の最大値)
  • right[ i ] := 右から i 文字分についての最適解 (o の個数の最大値)

という風にしてあげて、これを前処理で求めてあげればよさそうである。これだけだと一瞬不安になる。left[ i ] + right[ N - i - 1 ] = K であったとしても、left の解と right の解とが conflict してしまう場合があるように思えてしまうのだ (そのような i も採用しなければならなそうだ)。

でもそんなことは起こらない。なぜならそのような状況になる場合、たとえ S[ i ] = o を潰さなかったとしても、全体の最適解は K にならないのだ。

計算量は  O(N)

#include <bits/stdc++.h>
using namespace std;

int N, K, C;
vector<int> sub(const string &S) {
    int N = S.size();
    int cur = 0, last = -C-1;
    vector<int> res(N+1, 0);
    for (int i = 0; i < N; ++i) {
        if (i - last > C && S[i] == 'o') ++cur, last = i;
        res[i+1] = cur;
    }
    return res;
}

void solve(const string &S) {
    const auto &left = sub(S);
    string T = S;
    reverse(T.begin(), T.end());
    const auto &right = sub(T);
    for (int i = 0; i < N; ++i) {
        if (S[i] == 'x') continue;
        if (left[i] + right[N-i-1] < K) cout << i+1 << endl;
    }
}

int main() {
    string S;
    cin >> N >> K >> C >> S;
    solve(S);
}