ARC101 D. Median of Medians
問題
長さNの数列aが与えられる。
数列aのすべての連続部分列について中央値を並べ、新たに数列mを作る。
この数列mの中央値を求めよ。
ただし、偶数個の数列の中央値は、ソートし小さい方からfloor(M/2)+1番目の要素とする。( (10, 20, 30, 40)の場合、中央値は30)
制約
- 1 <= N <= 10^5
- 1 <= a_i <= 10^9
- a_iは整数
解法
求めるのは、数列mにおける中央値だが、数列mの要素数はN*(N+1)/2個もあるため、線形解法は通らない。
そこで、これよりも高速な二分探索で求めることを考える。
二分探索で中央値を求めることを考えると、「ある値X以上の数列mの要素数」がわかれば、その要素数が数列mの要素の半分を超えるところを求めればよい、ことがわかる。
次に、「ある値X以上の数列mの要素数」、すなわち、「数列aでの連続部分列で中央値がX以上の個数」を高速で求めたい。
ある連続部分列で中央値がX以上になるのは、連続部分列内の要素について「X未満の要素数 <= X以上の要素数」になる場合である。
これは以下のように考えることでO(NlogN)で求めることができる。
数列aにおいて、「X未満の要素」を-1、「X以上の要素」を+1とした数列bを考える。
数列a=(10,30,20)でX=20なら、数列b=(-1,+1,+1)となる。
求めたい「X未満の要素数 <= X以上の要素数となっている連続部分列の個数」は、数列bにおいて「総和が0以上になっている連続部分列の個数」になる。
連続部分列の総和が0以上かどうかは、累積和(最初に0を加えたN+1要素)を取って、lとrの位置での累積和を見て、「位置lでの累積和<=位置rでの累積和」になっていればよい。
ここで、累積和の値である2点の大小関係が異なる個数というのは「転倒数」であり、BITなどでO(NlogN)で求められるので、連続部分列の個数から転倒数を引けば、求めたい連続部分列の個数が求まる。
