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

hamayanhamayan's blog

Robot Takahashi [日鉄ソリューションズプログラミングコンテスト2022(AtCoder Beginner Contest 257) C]

https://atcoder.jp/contests/abc257/tasks/abc257_c

前提知識

  • SegTree (だが、使わない解法もあると思う)

解説

https://atcoder.jp/contests/abc257/submissions/32751195

全探索解法に向けて

まずは全探索解法がないかを模索してみる。
Xが実数全体を動くとされているが、実数空間を全探索することはむずかしい。
よりXの値を分かりやすく設定でき、かつ、意味のあるものにできないか?

体重が{2,5,9}をサンプルとして考えてみる。
まず、X=1とすると、全ての人は大人であると期待されることになる。
境界線を|と書くことにすると|259のような状況を考えていることになる。
X=1の次にX=1.5を考えてもよさそうだが、これは状況的には|259と全く変わりがない。
次に状況が変わるのは、Xが2を上回ったときである。
例えばX=3の時は2|59となるので、これは違う状態としてチェックが必要である。
次はX=4もよさそうだが、これも状況が変わらないので次はX=6あたりで25|9をチェックすることにする。
つまり、全探索するべきものはXというより「境界線」であり、
境界線は数の間と前後に置くことができるのでN+1通りとなる。
具体的に{2,5,9}のサンプルでいうと|259 2|59 25|9 259|の4通り。

まとめると、実数全体では全探索のイメージがつかめないので、境界線を全探索する。

全探索解法でやること

順番が答えに影響することはないので、体重についてはソートしておこう。
一旦、体重が同じ人がいないものと考えていこう。
全探索で境界線をおくと、正直具体的な体重については特に興味がなくなるので、
|010 0|10 01|0 010|
みたいな状態に対して境界線より前に0が何個あるか、後ろに1が何個あるかを求めて和を取ればf(X)になる。

このように「とある境界より前にある0の個数」「とある境界より後ろにある1の個数」が高速に求められれば、
境界線を全探索しても解けそうだ。

個数を数えるには

この部分は色々やり方が存在する。
一番計算量がよさそうなのは累積和を使用する方法だし、
境界線を全探索しながら差分だけ計算していくやり方もあるだろう。

一旦計算量は悪いが、データ構造を素直に適用するSegTree解法を紹介しておく。
Segtreeはとある配列に対して区間の総和などが取れるデータ構造である。
AtCoder Libraryにもある。
(説明は結構大変なので、SegTreeを知らない場合は一旦問題を止めてググってみてください…)
0用と1用の2つのSegTreeを用意しよう。

zero[L,R] = 区間[L,R)にある0の個数を返す
one[L,R] = 区間[L,R)にある1の個数を返す

詳しい実装はソースコードを見てもらう方が分かりやすいと思う。
zeroについては0がある場所に1を置いておけば区間総和すれば、区間にある0の個数が返ってくる。
これを使えば、zero[0,x) + one[x,N)みたいな感じでf(X)を求めることができる。
あとは、それの最大値を答えればAC。

体重が同じものはどうするか?

極端な例として、体重が同じ3名がいて、001みたいな大人子供設定だったとする。
体重で単純にソートして、大人子供の01でもソートすると001みたいな感じになる。
この状態で全探索をすると以下のようになる。
|001 0|01 00|1 001|
こうすると3番目で一番良い結果が出て3が答えになるように見えるが、
実際は全員同じ体重なので境界線が中途半端な所に来ることはなく、
|001 001|
の2択しかありえない。
よって単純にソートしてしまうと問題が発生する。

ちゃんとやるなら、体重が同じ人をまとめてしまって、
上のsegtreeでは1ではなくて人数を入れるのがいいと思うが、自分はソートを工夫することで切り抜けた。

体重で昇順ソートをして、大人子供設定は降順ソートすることにする。
この場合でも状態として
|100 1|00 10|0 100|
が発生してしまうが1と0が反転しているおかげで、中途半端に境界線が入っている状態がそうでないものよりも
よいf(X)になってしまうことを防いでいる。
この状態ではより良い結果を出すためには境界線を先頭か末尾に入れるしかなくなるので、
最大値を最終的に取ったときに中途半端な状態は無視される。
このような工夫で体重の重複を乗り切った。

int N;
string S;
int W[201010];
vector<pair<int, char>> WS;

int op(int a, int b) { return a + b; }
int e() { return 0; }

void _main() {
    cin >> N >> S;
    rep(i, 0, N) cin >> W[i];
    rep(i, 0, N) WS.push_back({ W[i], -S[i] });
    sort(all(WS));

    segtree<int, op, e> zero(vector<int>(N, 0));
    segtree<int, op, e> one(vector<int>(N, 0));
    rep(i, 0, N) {
        if (-WS[i].second == '0') zero.set(i, 1);
        else one.set(i, 1);
    }

    int ans = -1;
    rep(i, 0, N + 1) chmax(ans, zero.prod(0, i) + one.prod(i, N));
    cout << ans << endl;
}