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::collections::BTreeSet;

fn main() {
    let mut sc = Scanner::new();
    let n: usize = sc.read();
    let k: i64 = sc.read();

    let a: Vec<i64> = sc.read_vec(n);
    let mut b = vec![0; n + 1];
    for i in 0..n {
        b[i + 1] = b[i] + a[i] - k;
    }

    let mut set = BTreeSet::new();
    for &b in &b {
        set.insert(b);
    }
    let set: Vec<i64> = set.iter().map(|&b| b).collect();
    let mut c = vec![0; n + 1];
    for i in 0..(n + 1) {
        c[i] = set.binary_search(&b[i]).unwrap();
    }

    let mut bit = FenwickTree::new(n + 1);

    let mut ans = 0;
    bit.add(c[0], 1);
    for r in 1..(n + 1) {
        ans += bit.sum_one(c[r] + 1);
        bit.add(c[r], 1);
    }
    println!("{}", ans);
}

pub struct FenwickTree {
    n: usize,
    data: Vec<usize>,
}

impl FenwickTree {
    /// Constructs a new `FenwickTree`. The size of `FenwickTree` should be specified by `size`.
    pub fn new(size: usize) -> FenwickTree {
        FenwickTree {
            n: size + 1,
            data: vec![0; size + 1],
        }
    }

    fn add(&mut self, k: usize, value: usize) {
        let mut x = k;
        while x < self.n {
            self.data[x] += value;
            x |= x + 1;
        }
    }

    /// Returns a sum of range `[l, r)`
    pub fn sum(&self, l: usize, r: usize) -> usize {
        return self.sum_one(r) - self.sum_one(l);
    }

    /// Returns a sum of range `[0, k)`
    pub fn sum_one(&self, k: usize) -> usize {
        if k >= self.n {
            panic!("");
        }

        let mut result = 0;
        let mut x = k as i32 - 1;
        while x >= 0 {
            result += self.data[x as usize];
            x = (x & (x + 1)) - 1;
        }

        return result;
    }
}

struct Scanner {
    ptr: usize,
    length: usize,
    buf: Vec<u8>,
    small_cache: Vec<u8>,
}

#[allow(dead_code)]
impl Scanner {
    fn new() -> Scanner {
        Scanner {
            ptr: 0,
            length: 0,
            buf: vec![0; 1024],
            small_cache: vec![0; 1024],
        }
    }

    fn load(&mut self) {
        use std::io::Read;
        let mut s = std::io::stdin();
        self.length = s.read(&mut self.buf).unwrap();
    }

    fn byte(&mut self) -> u8 {
        if self.ptr >= self.length {
            self.ptr = 0;
            self.load();
            if self.length == 0 {
                self.buf[0] = b'\n';
                self.length = 1;
            }
        }

        self.ptr += 1;
        return self.buf[self.ptr - 1];
    }

    fn is_space(b: u8) -> bool {
        b == b'\n' || b == b'\r' || b == b'\t' || b == b' '
    }

    fn read_vec<T>(&mut self, n: usize) -> Vec<T>
    where
        T: std::str::FromStr,
        T::Err: std::fmt::Debug,
    {
        (0..n).map(|_| self.read()).collect()
    }

    fn read<T>(&mut self) -> T
    where
        T: std::str::FromStr,
        T::Err: std::fmt::Debug,
    {
        let mut b = self.byte();
        while Scanner::is_space(b) {
            b = self.byte();
        }

        for pos in 0..self.small_cache.len() {
            self.small_cache[pos] = b;
            b = self.byte();
            if Scanner::is_space(b) {
                return String::from_utf8_lossy(&self.small_cache[0..(pos + 1)])
                    .parse()
                    .unwrap();
            }
        }

        let mut v = self.small_cache.clone();
        while !Scanner::is_space(b) {
            v.push(b);
            b = self.byte();
        }
        return String::from_utf8_lossy(&v).parse().unwrap();
    }
}