AtCoder Beginner Contest 396 F

https://atcoder.jp/contests/abc396/tasks/abc396_f

まず初期の転倒数を数えます。入力例2

5 6
5 3 5 0 1

で、最小の0より左の数を数えます。このとき、

1 1 1 1 1

という配列を用意しておいて、Aの0の位置より左の和を数えます。そして、その位置を0にします。

1 1 1 0 1

次に小さいのは1なので、その位置より左の和は3で、というように和を取って行きます。これはBITで簡単に実現できます。

そのあとは1を全部に足すのですが、全部を動かすのは面倒なので5を-1とすると考えます。そうしたときの転倒数の増減を数えます。右の5はその右に2つあるので-2、左に2つあるけど5が1つあるので+1、左の5はその右に4つあるけど、5があるので-3で、トータルで-4です。結局、5が重複するのは相殺するので、ポジションをpとすると増減は、

となります。

// Rotated Inversions
#![allow(non_snake_case)]


//////////////////// library ////////////////////

fn read<T: std::str::FromStr>() -> T {
    let mut line = String::new();
    std::io::stdin().read_line(&mut line).ok();
    line.trim().parse().ok().unwrap()
}

fn read_vec<T: std::str::FromStr>() -> Vec<T> {
    read::<String>().split_whitespace()
            .map(|e| e.parse().ok().unwrap()).collect()
}


//////////////////// Binary Indexed Tree ////////////////////

type Value = i64;

struct BIT {
    v: Vec<Value>
}

impl BIT {
    fn ceil_two_pow(n: usize) -> usize {
        if n == 1 { 1 } else { BIT::ceil_two_pow((n+1)/2) * 2 }
    }
    
    fn new(m: usize) -> BIT {
        let n = BIT::ceil_two_pow(m);
        let v: Vec<Value> = vec![0; n+1];
        BIT { v }
    }
    
    // v_1 + ... + v_n
    fn sum(&self, n: usize) -> Value {
        let mut s: Value = 0;
        let mut m: Value = n as Value;
        while m != 0 {
            s += self.v[m as usize];
            m -= m & (-m)
        }
        s
    }
    
    fn add(&mut self, n: usize, e: Value) {
        let mut m: Value = n as Value;
        while (m as usize) < self.v.len() {
            self.v[m as usize] += e;
            m += m & (-m)
        }
    }
}


//////////////////// process ////////////////////

fn read_input() -> (i64, Vec<i64>) {
    let v: Vec<i64> = read_vec();
    let M = v[1];
    let A: Vec<i64> = read_vec();
    (M, A)
}

fn calc_initial_inversion_number(v: &Vec<(i64, usize)>) -> i64 {
    let N = v.len();
    let mut tree = BIT::new(N);
    for i in 0..N {
        tree.add(i+1, 1)
    }
    
    let mut s: i64 = 0;
    for &(_, p) in v.iter() {
        s += tree.sum(p);
        tree.add(p+1, -1)
    }
    s
}

fn F(M: i64, A: Vec<i64>) {
    let mut v: Vec<(i64, usize)> = A.iter().enumerate().
                                    map(|(p, &a)| (a, p)).collect();
    v.sort();
    
    let mut s = calc_initial_inversion_number(&v);
    println!("{}", s);
    
    let N = A.len();
    for n in (1..M).rev() {
        loop {
            if let Some(&(a, p)) = v.last() {
                if a == n {
                    s += (p + p) as i64;
                    s -= (N - 1) as i64;
                    v.pop();
                    continue
                }
            }
            break
        }
        println!("{}", s)
    }
}

fn main() {
    let (M, A) = read_input();
    F(M, A)
}