やったこと

問題を解く
「愚直にやると $O(N^{2}\log N)$ になってTLEしますね」
「想定解は $O(N^{2})$ より低いです」
「まじで」
数時間睨むと $O(N^{5/3}\log N)$ 解法が生える
1secくらいで通る
writer解と解法もorderも違うことが判明する
large2_4.txt, large2_5.txtを追加する
- large2_4.txt: $K=N$ (奇数) のときに最大値となるケース
- large2_5.txt: $K=N$ (偶数) のときに最大値となるケース
writer解がTLEする (multiset::countのオーダー問題)
(前日夜に) Javaでもなんとか通ることを確認する (2.8sec)
$A_i$ は整数という制約を入れる (忘れてた)
$K=1 \ldots 400$ と $K=N$ だけ調べればACすることがわかったため、急遽テストケースを追加 (large2_6.txt, large2_7.txt)
当日夕方にJava解がgeneratorによるテストケースに落とされる
JavaのPriorityQueue#eraseが線形であると気づく
それでも3秒切ってるしもう時間ないしまぁいいか

自分の解法

愚直解: $O(N^{2}\log N)$

まずは、要素数 $N$ に対して $O(\log N)$ で要素を追加および中央値の取得ができるデータ構造が必要となります。これは、C++のstd::priority_queueやstd::multisetを2個使うことでできます。RBST(Randomized Binary Search Tree)のようなデータ構造で殴っても大丈夫です。このようなデータ構造があれば、 $K=1 \ldots N$ で左右から区間中央値の累積和の最大値を取っていくことで $O(N^{2} \log N)$ で解けますが、TLに間に合いません。

Tester解: $O(N^{5/3}\log N)$

たとえば、 $K=a$ と $K=a+1$ の場合を考えてみましょう。すると、 $a$ が大きい場合は上記のようなデータ構造を毎回作り直すよりもスライドさせたほうが計算量が小さくなりそうです。スライドするためには $O(\log N)$ で要素を削除する必要があり、これはstd::priority_queueではできないですがstd::multisetだとできます。 $K=a$ から $K=a+1$ にスライドさせる場合、データ構造の数を $r = \lfloor N/a \rfloor$ とおくと要素追加および削除の計算量は $O(r^{2} \log a)$ になります。ここで、 $K=1 \ldots N^{2/3}$ は愚直解の方法で、 $K=N^{2/3} \ldots N$ はスライドさせて解くと、前半部分の計算量は $O(N^{5/3}\log N)$ 、後半部分も $r \le N^{1/3}$ より $O((N - N^{2/3}) r^{2} \log N) = O(N^{5/3}\log N)$ となり、 $N=10000$ でもぎりぎり間に合います。

C++の提出コード https://yukicoder.me/submissions/382387

また、この問題では中央値の取得でしたが、k番目の値を取得するためのデータ構造については以下の記事にまとまっています。 qiita.com

(追記)

コンテストの提出を見ていたらWavelet Matrixでより高速に解けることがわかりました。勉強になる。