ARC074 D. 3N Numbers
問題
長さが3*Nの数列vが与えられる。この数列からちょうどN個の要素を取り除き、「前からN個の配列a」と「後ろからN個の配列b」を考える。
「aの総和-bの総和」の最大値を求めよ。
制約
1 <= N <= 10^5
1 <= 要素の値 <= 10^9
解説
数列vにおいて適当に1か所場所を決めると、そこより前からの要素が配列a、それ以降の要素が配列bと割り当てて考えることができる。
最大値がほしいので、配列aの方は大きい値からN個、配列bの方は小さい値からN個を取ればよい。
愚直に計算してしまうとO(N^2)でTLEしてしまうので、高速に計算できないかを考える。
欲しいのは「集合Sに要素xを追加したとき値の大きいor小さい上位N個の和」が高速に求められるようなもの。
上位N個を保持するのに使えるデータ構造はheap(priority_queue)で、これを使って、
- 和と上位N個の要素を保持しておく
- 要素xを追加したら和とheapに入れ、heapから1つ要素を取り出し、和から引いて捨てる
のようにする。
値の追加はO(logN)、上位N個の和の取得はO(1)で求められ、十分に速い。
配列bの方は、小さい方からN個&後ろから前に向かって考えると配列aと同様に高速に求められる。
別途配列などに配列aと配列bの上位N件の和の値を入れておくなどして、最後に線形に「位置iより前の大きい上位N個-位置i以降の小さい上位N個」の最大値を求めればO(NlogN)で答えが求まる。
反省
最初、何も考えずに、位置に対して最大値は"上に凸"になるのでは?と思って確かめずに投げてWAやTLEを出しまくってしまった。
ランダムに生成させてみると、
3 10 5 2 4 2 18 10 10 7
のとき、最大値は各位置で「x,x,4,0,-8,6,x,x,x」のように簡単にそうならないケースが見つかる。
ちゃんと確かめてからコードを書き始めること。