Problem:
Score : 600 points

Problem Statement

You are given an integer sequence of length N, a = {a_1, a_2, …, a_N}, and an integer K.

a has N(N+1)/2 non-empty contiguous subsequences, {a_l, a_{l+1}, …, a_r} (1 ≤ l ≤ r ≤ N). Among them, how many have an arithmetic mean that is greater than or equal to K?

Constraints

All input values are integers.

1 ≤ N ≤ 2 \times 10^5

1 ≤ K ≤ 10^9

1 ≤ a_i ≤ 10^9

Input

Input is given from Standard Input in the following format:

N K
a_1
a_2
:
a_N

Output

Print the number of the non-empty contiguous subsequences with an arithmetic mean that is greater than or equal to K.

Sample Input 1

3 6
7
5
7

Sample Output 1

5

All the non-empty contiguous subsequences of a are listed below:

{a_1} = {7}

{a_1, a_2} = {7, 5}

{a_1, a_2, a_3} = {7, 5, 7}

{a_2} = {5}

{a_2, a_3} = {5, 7}

{a_3} = {7}

Their means are 7, 6, 19/3, 5, 6 and 7, respectively, and five among them are 6 or greater. Note that {a_1} and {a_3} are indistinguishable by the values of their elements, but we count them individually.

Sample Input 2

1 2
1

Sample Output 2

0

Sample Input 3

7 26
10
20
30
40
30
20
10

Sample Output 3

13

Solution in Rust (1.15.1) (judge verdict: Accepted)

#[allow(unused_imports)]
use std::cmp::*;
#[allow(unused_imports)]
use std::collections::*;

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

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

fn main() {
    let mut v: Vec<i64> = read_vec();
    let n = v[0] as usize;
    let k = v[1];
    let mut a = vec![0; n];
    for i in 0..n {
        a[i] = read::<i64>();
    }
    let a = a.iter().map(|&num| num - k).collect::<Vec<_>>();
    let mut sums = vec![-k; n + 1];
    for i in 1..n + 1 {
        sums[i] = sums[i - 1] + a[i - 1];
    }

    let indexes = argsort(&sums).iter().map(|&x| x).collect::<Vec<_>>();
    let mut c = vec![0; n + 1];
    for i in 0..n + 1 {
        c[indexes[i]] = n + 1 - i;
    }

    let mut bt = BinaryIndexTree::new(indexes.len());
    let mut ans = 0;
    for i in 0..indexes.len() {
        ans += i - bt.sum(c[i]) as usize;
        bt.add(c[i] as usize, 1);
    }
    println!("{:?}", ans);
}

fn argsort<T: std::cmp::Ord + std::marker::Copy>(v: &Vec<T>) -> Vec<usize> {
    let mut sort_v = Vec::new();
    for i in 0..v.len() {
        sort_v.push((v[i], i));
    }
    sort_v.sort();
    sort_v.iter().map(|x| x.1).collect::<Vec<usize>>()
}

struct BinaryIndexTree {
    bit: Vec<i64>,
    n: usize,
}

impl BinaryIndexTree {
    fn new(n: usize) -> BinaryIndexTree {
        BinaryIndexTree {
            bit: vec![0; n + 1],
            n: n,
        }
    }

    fn sum(&self, i: usize) -> i64 {
        assert!(i > 0);
        let mut i = i;
        let mut s = 0i64;
        while i > 0 {
            s += self.bit[i];
            i -= (i as i64 & -(i as i64)) as usize;
        }
        s
    }

    fn add(&mut self, i: usize, x: i64) {
        assert!(i > 0);
        let mut i = i;
        while i <= self.n {
            self.bit[i as usize] += x;
            i += (i as i64 & -(i as i64)) as usize;
        }
    }
}