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)

use std::io::{self, Stdin};
use std::str::{self, FromStr};
use std::error::Error;

fn exec() {
    let mut sc = Scanner::new();
    let n: usize = sc.ne();
    let k: i64 = sc.ne();
    let a: Vec<i64> = (0..n).map(|_| sc.ne()).collect();
    let mut b = Vec::new();
    let mut s = vec![0; n + 1];
    for i in 0..n {
        s[i + 1] = s[i] + a[i];
        s[i + 1] -= k;
        b.push((s[i + 1], i + 1));
    }
    b.sort();
    let mut mp = vec![0; n + 1];
    for i in 0..n {
        mp[b[i].1] = i;
    }
    let mut fwt = FenwickTree::new(n + 2, 0i64);
    let mut sum = 0;
    for i in (0..n).rev() {
        fwt.add(mp[i + 1], 1);
        let bound = lower_bound(&b, &(s[i], 0));
        sum += fwt.sum(n + 2) - fwt.sum(bound);
    }
    println!("{}", sum);
}

fn lower_bound<T: Ord>(a: &[T], q: &T) -> usize {
    let (mut lb, mut ub) = (-1, a.len() as i64);
    while ub - lb > 1 {
        let mid = (lb + ub) / 2;
        if a[mid as usize] >= *q {
            ub = mid;
        } else {
            lb = mid;
        }
    }
    ub as usize
}

#[allow(dead_code)]
struct FenwickTree<T> {
    zero: T,
    n: usize,
    data: Vec<T>,
}
use std::ops::AddAssign;
#[allow(dead_code)]
impl<T: Copy + AddAssign> FenwickTree<T> {
    fn new(n_: usize, z_: T) -> FenwickTree<T> {
        FenwickTree {
            zero: z_,
            n: n_,
            data: vec![z_; n_],
        }
    }
    //sum of [0, j)
    fn sum(&self, j: usize) -> T {
        let mut res = self.zero;
        let mut x = j as i64 - 1;
        while x >= 0 {
            res += self.data[x as usize];
            x = (x & (x + 1)) - 1;
        }
        res
    }
    // add w to a[j]
    fn add(&mut self, j: usize, w: T) {
        let mut x = j;
        while x < self.n {
            self.data[x] += w;
            x |= x + 1;
        }
    }
}

fn main() {
    const STACK: usize = 16 * 1024 * 1024;
    let _ = std::thread::Builder::new()
        .stack_size(STACK)
        .spawn(|| { exec(); })
        .unwrap()
        .join()
        .unwrap();
}

#[allow(dead_code)]
struct Scanner {
    stdin: Stdin,
    id: usize,
    buf: Vec<u8>,
}

#[allow(dead_code)]
impl Scanner {
    fn new() -> Scanner {
        Scanner {
            stdin: io::stdin(),
            id: 0,
            buf: Vec::new(),
        }
    }
    fn next_line(&mut self) -> Option<String> {
        let mut res = String::new();
        match self.stdin.read_line(&mut res) {
            Ok(0) => None,
            Ok(_) => Some(res),
            Err(why) => panic!("error in read_line: {}", why.description()),
        }
    }
    fn next<T: FromStr>(&mut self) -> Option<T> {
        while self.buf.len() == 0 {
            self.buf = match self.next_line() {
                Some(r) => {
                    self.id = 0;
                    r.trim().as_bytes().to_owned()
                }
                None => return None,
            };
        }
        let l = self.id;
        assert!(self.buf[l] != b' ');
        let n = self.buf.len();
        let mut r = l;
        while r < n && self.buf[r] != b' ' {
            r += 1;
        }
        let res = match str::from_utf8(&self.buf[l..r])
                  .ok()
                  .unwrap()
                  .parse::<T>() {
            Ok(s) => Some(s),
            Err(_) => {
                panic!("parse error, {:?}",
                       String::from_utf8(self.buf[l..r].to_owned()))
            }
        };
        while r < n && self.buf[r] == b' ' {
            r += 1;
        }
        if r == n {
            self.buf.clear();
        } else {
            self.id = r;
        }
        res
    }
    fn ne<T: FromStr>(&mut self) -> T {
        self.next::<T>().unwrap()
    }
}