バイトの競プロメモ

主に競技プログラミング

D - Median of Medians AtCoder Regular Contest 101

D: Median of Medians - AtCoder Regular Contest 101 | AtCoder
問題概略
長さMの数列の任意の区間[lr]の中央値でなる数列の中央値を求めよ。
10^5 N
10^9 A

解法
xが中央値が満たす条件について考えてみる。
x以上の要素がs/2以上ある。
そのような物の内最大である。

上のように気づけると二分探索しようと考えることが出来る。

つまり、任意の区間について中央値がx以上となるものが半分以上あれば条件を満たす。
ある区間[lr]の中央値がx以上であることを調べるには、[]をx以上なら1,その他は-1として、合計が0以上になれば良い。
これはSr-Sl>=0を意味するが、これをも数えるには累積和を取ってやり、任意のrについて、Sl <= Srとなるようなlを数えればよく。
反転数を数えるときと同様にBITでとける。
今回はあくまで累積和なので、先頭に番兵を追加してやらないと左端から自分までの合計の時の反転数を数えられないことに注意

public static void solve() throws Exception {
        //longを忘れるなオーバーフローするぞ
        N = ni();
        A = nia(N);
        SOA = A.clone();
        Arrays.sort(SOA);
        System.out.println(binarysearch(0, N));
    }
 
    static class BIT {
        long[] dat;
        int N;
 
        BIT(int n) {
            N = n;
            dat = new long[n + 1];
        }
 
        // 0 index
        void add(int i, int v) {
            if (i >= dat.length) return;
            dat[i] += v;
            add(i + (i & -i), v);
        }
 
        //開区間
        long _sum(int i) {
            if (i <= 0) return 0;
            return dat[i] + _sum(i - (i & -i));
        }
 
        //半開区間
        long sum(int l, int r) {
            return _sum(r) - _sum(l);
        }
    }
 
    public static boolean calc(long v) {
        int va = SOA[(int) v];
        int[] rui = new int[N + 1];
        for (int i = 0; i < N; i++) {
            if (A[i] >= va) rui[i + 1]++;
            else rui[i + 1]--;
        }
        for (int i = 0; i < N; i++) {
            rui[i + 1] += rui[i];
        }
        int key1 = 114514;
        int key2 = 215415;
        for (int i = 0; i < N+1; i++) {
            rui[i] += key1;
        }
        BIT b = new BIT(N + key2);
        long res = 0;
        for (int i = 0; i < N + 1; i++) {
            res += b._sum(rui[i]);
            b.add(rui[i], 1);
        }
        return N * 1L * (N + 1) / 2 <= res*2;
    }
 
    //条件を満たす最大値、あるいは最小値を求める
    static long binarysearch(long ok, long ng) {
        //int ok = 0; //解が存在する
        //int ng = N; //解が存在しない
        while (Math.abs(ok - ng) > 1) {
            long mid;
            if (ok < 0 && ng > 0 || ok > 0 && ng < 0) mid = (ok + ng) / 2;
            else mid = ok + (ng - ok) / 2;
 
            if (calc(mid)) {
                ok = mid;
            } else {
                ng = mid;
            }
        }
        return SOA[(int) ok];
    }