Problem:
Score : 500 points

Problem Statement

Given are a sequence of N positive integers A_1, A_2, \ldots, A_N, and a positive integer K.

Find the number of non-empty contiguous subsequences in A such that the remainder when dividing the sum of its elements by K is equal to the number of its elements. We consider two subsequences different if they are taken from different positions, even if they are equal sequences.

Constraints

All values in input are integers.

1 \leq N \leq 2\times 10^5

1 \leq K \leq 10^9

1 \leq A_i \leq 10^9

Input

Input is given from Standard Input in the following format:

N K
A_1 A_2 \cdots A_N

Output

Print the number of subsequences that satisfy the condition.

Sample Input 1

5 4
1 4 2 3 5

Sample Output 1

4

Four sequences satisfy the condition: (1), (4,2), (1,4,2), and (5).

Sample Input 2

8 4
4 2 4 2 4 2 4 2

Sample Output 2

7

(4,2) is counted four times, and (2,4) is counted three times.

Sample Input 3

10 7
14 15 92 65 35 89 79 32 38 46

Sample Output 3

8

Solution in Rust (1.15.1) (judge verdict: Time Limit Exceeded)

#![allow(unused_imports)]
use std::cmp::{max, min};
use std::collections::{BTreeMap, BTreeSet, BinaryHeap, HashMap, HashSet, VecDeque};

use nyan::*;

fn main() {
    let stdin = std::io::stdin();
    let mut sc = Scanner::new(stdin.lock());

    let n: usize = sc.read();
    let k: u64 = sc.read();
    let a: Vec<u64> = sc.read_vec(n);

    let sums = a
        .iter()
        .scan(0, |state, &ai| {
            *state += ai;
            Some(*state)
        })
        .collect::<Vec<_>>();

    let mut ans = 0;
    for i in 0..n {
        for j in i..n {
            let s = sums[j] - if i == 0 { 0 } else { sums[i - 1] };
            let l = j - i + 1;
            if s % k == l as u64 {
                ans += 1;
            }
        }
    }
    println!("{}", ans);
}

#[allow(dead_code)]
mod nyan {
    use std::cmp::*;
    use std::fmt::{Debug, Display};
    use std::io::{self, BufRead};
    use std::iter::FromIterator;
    use std::str::FromStr;

    // https://doc.rust-lang.org/1.39.0/src/core/num/mod.rs.html#4473-4523
    #[inline]
    pub fn is_ascii_whitespace(c: char) -> bool {
        match c {
            '\t' | '\n' | '\x0C' | '\r' | ' ' => true,
            _ => false,
        }
    }

    #[inline]
    pub fn is_ascii_linebreak(c: char) -> bool {
        match c {
            '\n' | '\r' => true,
            _ => false,
        }
    }

    struct Reader<R> {
        reader: R,
        buffer: String,
        cursor: usize,
    }

    impl<R: BufRead> Reader<R> {
        fn new(reader: R) -> Reader<R> {
            Reader {
                reader: reader,
                buffer: String::with_capacity(32),
                cursor: 0,
            }
        }

        #[inline]
        fn trim_start(&mut self) {
            match self.buffer[self.cursor..].find(|c| !is_ascii_whitespace(c)) {
                Some(l) => self.cursor += l,
                None => self.cursor = self.buffer.len(),
            }
        }

        fn prepare_buffer(&mut self) -> io::Result<()> {
            self.trim_start();

            while self.buffer.len() - self.cursor == 0 {
                self.buffer.clear();
                match self.reader.read_line(&mut self.buffer) {
                    Err(e) => return Err(e),
                    Ok(0) => return Err(io::Error::from(io::ErrorKind::UnexpectedEof)),
                    Ok(_) => {
                        self.cursor = 0;
                        self.trim_start();
                    }
                }
            }

            Ok(())
        }

        fn next_word(&mut self) -> io::Result<&str> {
            self.prepare_buffer()?;

            match self.buffer[self.cursor..].find(is_ascii_whitespace) {
                Some(l) => {
                    let start = self.cursor;
                    self.cursor += l;
                    Ok(&self.buffer[start..self.cursor])
                }
                None => Err(io::Error::from(io::ErrorKind::Other)),
            }
        }

        fn next_line(&mut self) -> io::Result<&str> {
            self.prepare_buffer()?;

            match self.buffer[self.cursor..].find(is_ascii_linebreak) {
                Some(l) => {
                    let start = self.cursor;
                    self.cursor += l;
                    Ok(&self.buffer[start..self.cursor])
                }
                None => Err(io::Error::from(io::ErrorKind::Other)),
            }
        }
    }

    pub struct Scanner<R> {
        reader: Reader<R>,
    }

    impl<R: BufRead> Scanner<R> {
        pub fn new(reader: R) -> Scanner<R> {
            Scanner {
                reader: Reader::new(reader),
            }
        }

        pub fn read<T>(&mut self) -> T
        where
            T: FromStr,
            <T as FromStr>::Err: Debug,
        {
            self.reader.next_word().unwrap().parse().unwrap()
        }

        pub fn read_vec<T, U>(&mut self, n: usize) -> T
        where
            T: FromIterator<U>,
            U: FromStr,
            <U as FromStr>::Err: Debug,
        {
            (0..n).map(|_| self.read()).collect()
        }

        pub fn read_vec2d<T, U, V>(&mut self, m: usize, n: usize) -> T
        where
            T: FromIterator<U>,
            U: FromIterator<V>,
            V: FromStr,
            <V as FromStr>::Err: Debug,
        {
            (0..m).map(|_| self.read_vec(n)).collect()
        }

        pub fn read_line<T: FromIterator<char>>(&mut self) -> T {
            self.reader.next_line().unwrap().chars().collect()
        }

        pub fn read_lines<T, U>(&mut self, n: usize) -> T
        where
            T: FromIterator<U>,
            U: FromIterator<char>,
        {
            (0..n).map(|_| self.read_line()).collect()
        }
    }

    pub fn print_vec<T: Display>(v: &[T]) {
        if v.is_empty() {
            return;
        }

        print!("{}", v[0]);
        for i in v.iter().skip(1) {
            print!(" {}", i);
        }
        println!();
    }

    // C++ like binary search functions (upper_bound, lower_bound, equal_range)
    pub trait BinarySearch<T> {
        fn lower_bound(&self, value: &T) -> usize;
        fn upper_bound(&self, value: &T) -> usize;
        fn equal_range(&self, value: &T) -> (usize, usize) {
            (self.lower_bound(value), self.upper_bound(value))
        }
    }

    impl<T: Ord> BinarySearch<T> for [T] {
        // https://en.cppreference.com/w/cpp/algorithm/lower_bound#Possible_implementation
        fn lower_bound(&self, value: &T) -> usize {
            let mut first = 0;
            let mut last = self.len();
            while last != first {
                let mid = (first + last) / 2;
                match self[mid].cmp(value) {
                    Ordering::Less => first = mid + 1,
                    _ => last = mid,
                }
            }
            first
        }

        // https://en.cppreference.com/w/cpp/algorithm/upper_bound#Possible_implementation
        fn upper_bound(&self, value: &T) -> usize {
            let mut first = 0;
            let mut last = self.len();
            while last != first {
                let mid = (first + last) / 2;
                match self[mid].cmp(value) {
                    Ordering::Greater => last = mid,
                    _ => first = mid + 1,
                }
            }
            first
        }
    }

    #[test]
    pub fn test_binary_search() {
        let data = vec![1, 1, 2, 3, 3, 3, 3, 4, 4, 4, 5, 5, 6];
        assert_eq!(data.lower_bound(&4), 7);
        assert_eq!(data.upper_bound(&4), 10);
    }

    // https://doc.rust-lang.org/std/cmp/struct.Reverse.html
    #[derive(PartialEq, Eq, Debug)]
    pub struct Rev<T>(pub T);

    impl<T: PartialOrd> PartialOrd for Rev<T> {
        fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
            other.0.partial_cmp(&self.0)
        }
    }

    impl<T: Ord> Ord for Rev<T> {
        fn cmp(&self, other: &Self) -> Ordering {
            other.0.cmp(&self.0)
        }
    }

    // Extended Euclidean algorithm
    // ax + by = gcd(x, y)
    pub fn extended_gcd(a: u64, b: u64) -> (u64, i64, i64) {
        if b == 0 {
            (a, 1, 0)
        } else {
            let (m, y, x) = extended_gcd(b, a % b);
            (m, x, y - (a / b) as i64 * x)
        }
    }

    pub fn gcd(a: u64, b: u64) -> u64 {
        if b == 0 {
            a
        } else {
            gcd(b, a % b)
        }
    }

    pub fn lcm(a: u64, b: u64) -> u64 {
        (a * b) / gcd(a, b)
    }

    // x^n (mod p)
    pub fn pow_mod(mut x: u64, mut n: u64, p: u64) -> u64 {
        let mut r = 1;
        while n > 0 {
            if n & 1 > 0 {
                r *= x;
                r %= p;
            }
            x *= x;
            x %= p;
            n >>= 1;
        }
        r
    }

    #[test]
    fn test_pow_mod() {
        assert_eq!(pow_mod(4, 13, 497), 445);
        assert_eq!(pow_mod(2, 90, 13), 12);
        assert_eq!(pow_mod(2, 99999999, 147), 134);
    }

    // a^{p-1} = 1 (mod p) (Fermat's little theorem)
    // a^{p - 2} * a = 1 → a^{-1} = a^{p-2}
    pub fn inv_mod(n: u64, p: u64) -> u64 {
        pow_mod(n, p - 2, p)
    }

    // s * (s + 1) * .. * t (mod p)
    pub fn prod_mod(s: u64, t: u64, p: u64) -> u64 {
        (s..t + 1).fold(1, |acc, x| acc * x % p)
    }

    // nCk (mod p)
    pub fn comb_mod(n: u64, k: u64, p: u64) -> u64 {
        // nCk = n! / (k! * (n - k)!)
        //     = ((n-k+1) * (n-k+2) * ... * n) * inv(k!)
        prod_mod(n - k + 1, n, p) * inv_mod(prod_mod(1, k, p), p) % p
    }

    #[inline]
    pub fn minmax<T: Ord>(a: T, b: T) -> (T, T) {
        if a < b {
            (a, b)
        } else {
            (b, a)
        }
    }

    #[test]
    fn test_minmax() {
        assert_eq!(minmax(1, 3), (1, 3));
        assert_eq!(minmax(3, 1), (1, 3));
    }

    pub struct UnionFind {
        n: usize,
        root: Vec<usize>,
        rank: Vec<usize>,
        size: Vec<usize>,
    }

    impl UnionFind {
        pub fn new(n: usize) -> UnionFind {
            UnionFind {
                n: n,
                root: (0..n).collect(),
                rank: vec![0; n],
                size: vec![1; n],
            }
        }

        pub fn root(&mut self, x: usize) -> usize {
            let r = self.root[x];
            if r == x {
                r
            } else {
                self.root[x] = self.root(r);
                self.root[x]
            }
        }

        pub fn merge(&mut self, x: usize, y: usize) {
            let rx = self.root(x);
            let ry = self.root(y);
            if rx == ry {
                return;
            }

            if self.rank[rx] < self.rank[ry] {
                self.root[rx] = ry;
                self.size[ry] += self.size[rx];
            } else {
                self.root[ry] = rx;
                self.size[rx] += self.size[ry];
                if self.rank[rx] == self.rank[ry] {
                    self.rank[rx] += 1;
                }
            }
        }

        pub fn size(&mut self, x: usize) -> usize {
            let rx = self.root(x);
            self.size[rx]
        }
    }
}