問題

https://atcoder.jp/contests/arc101/tasks/arc101_b

定義

$\displaystyle{ b }$ を昇順にソートして得られる数列を $\displaystyle{ b′ }$ とする。このとき、 $\displaystyle{ b′ }$ の $\displaystyle{ \frac{M}{2}+1 }$ 番目の要素の値を、b の中央値とする。ここで、/ は小数点以下を切り捨てる除算である。

解法

$\displaystyle{ f(X) = }$ 求める中央値が $\displaystyle{X}$ 以下となるかどうか

という判定条件を考える。

この判定条件が初めてtrueになる境界を二分探索で求めていく。

与えられる数列 $\displaystyle{ a }$ の区間 $\displaystyle{ [l,r) }$ は、 $\displaystyle{ a }$ の長さを $\displaystyle{ M }$ とおくと、 $\displaystyle{ _{M+1}C_2 }$ 個ある。

例) $\displaystyle{ [30,10,20] }$ の区間は、 $\displaystyle{ [[30], [10], [20], [30, 10], [10, 20], [30, 10, 20]] }$ の $\displaystyle{ _{3+1}C_2 }$ = 6通り

「考えられうる区間の半数」以上、すなわち「 $\displaystyle{ \frac{_{M+1}C_2}{2} }$ 個以上の区間」の中央値が $\displaystyle{ X }$ 以下であれば、 $\displaystyle{ f(X) = true }$ となる。

「区間の中央値が $\displaystyle{ X }$ 以下がどうか」は、特に区間をsortせずとも、

$\displaystyle{ X }$ 以上の要素を1, $\displaystyle{ X }$ 未満の要素を-1に変換
区間の総和が0未満となるかどうか

で判定できる。

区間の総和は、累積和を使うと処理しやすい。

$\displaystyle{ S_0 }$ = 0
$\displaystyle{ S_1 }$ = $\displaystyle{ a_0 }$
$\displaystyle{ S_2 }$ = $\displaystyle{ a_0 + a_1 }$
$\displaystyle{ S_3 }$ = $\displaystyle{ a_0 + a_1 + a_2 }$

とおくと、区間 $\displaystyle{ [l, r) }$ の総和は $\displaystyle{ S_r - S_l }$ と求めることができる。

よって、問題は次のように言い換えることができる。

数列 $\displaystyle{ a }$ の要素のうち、 $\displaystyle{ x }$ 以上の要素を $\displaystyle{ 1 }$ 、 $\displaystyle{ x }$ 未満の要素を $\displaystyle{ -1 }$ へと変換した数列を $\displaystyle{ b }$ とおく。
$\displaystyle{ b }$ の累積和数列 $\displaystyle{ S }$ において、 $\displaystyle{ S_r - S_l \lt 0 }$ となる区間 $\displaystyle{ [l, r) }$ の数が $\displaystyle{ \frac{_{M+1}C_2}{2} }$ 以上であるような、最小の $\displaystyle{ x }$ を求めよ。

ここで、

$\displaystyle{ S }$ において、 $\displaystyle{ S_r - S_l \lt 0 }$ となる区間 $\displaystyle{ [l, r) }$ の数

は、 $\displaystyle{ S }$ の転倒数を求めていることとほぼ同義である。

ここの転倒数を求める処理で、愚直に

        for l in 0..n as usize {
            for r in l..n as usize {
                if S[r] < S[l] {
                    cnt += 1
                }
            }
        }

とすると、 $\displaystyle{ O(N^2) }$ となり、実行時間切れとなる。 Fenwick Tree等を用いて、 $\displaystyle{ O(Nlog(N)) }$ あたりまで計算量を下げることが求められる。

Fenwick Treeを用いた転倒数算出については別エントリに記載した。

// https://atcoder.jp/contests/arc101/tasks/arc101_b

use proconio::input;

// nearly exceeds 10^9
const A_MAX: i64 = 1<<30;

struct Fenwick {
    tree: Vec<i64>,
}

impl Fenwick {
    /// create Fenwick Tree by given length
    pub fn new(len: usize) -> Self {
        Self { tree: vec![0; len] }
    }

    /// add `w` to index `i`
    /// `i` is 1-indexed
    pub fn add(&mut self, mut i: usize, w: i64) {
        while i < self.tree.len() {
            self.tree[i] += w;
            i += lsb(i);
        }
    }

    /// sum `[1, a]`
    /// `a` is 1-indexed
    pub fn sum(&self, mut i: usize) -> i64 {
        let mut res = 0;
        while i > 0 {
            res += self.tree[i];
            i -= lsb(i);
        }
        res
    }

    /// sum `[a, b)`
    /// a and b are 1-indexed
    pub fn partial_sum(&self, a: usize, b: usize) -> i64 {
        self.sum(b-1) - self.sum(a-1)
    }
}

/// return least significant bit of i
fn lsb(i: usize) -> usize {
    if i == 0 {
        0
    } else {
        i & !(i - 1)
    }
}

fn main() {
    input! {
        n: i64,
        mut a: [i64; n]
    }

    if n == 1 {
        println!("{}", a[0]);
        return
    }

    let mut left = 0;
    let mut right = A_MAX;
    let geta = n + 1;
    while right - left > 1 {
        let mid = (right + left) / 2;
        let mut inversion_number = 0;
        let mut ft = Fenwick::new(n as usize * 2 + 10);
        let mut sum = 0;
        ft.add((sum + geta) as usize, 1);
        for i in 0..n as usize {
            let val = if a[i] <= mid { 1 } else { -1 };
            sum += val;
            inversion_number += ft.partial_sum(1, (sum+geta) as usize);
            ft.add((sum+geta) as usize, 1);
        }

        if inversion_number as i64 > (n+1)*n/2/2 {
            right = mid;
        } else {
            left = mid;
        }
    }
    println!("{}", right);
}

2sec

2秒間でライムを刻め

ARC101 D Median of Medians

問題

定義

解法

references