AtCoder Beginner Contest 288 D

https://atcoder.jp/contests/abc288/tasks/abc288_d

 X_i, X_{i+1}, \cdots , X_{i+K-1}にcを加算するとはどういうことでしょう。Aを多項式、すなわち A_1 + A_2x + \cdots + A_Nx^{N-1}と考えると、 c(x^{i-1}+ \cdots x^{i+K-2})を足すことです。つまり、この問題は、 1 + x+ \cdots + x^{K-1}の倍数を足して0になるか、言い換えると多項式 1 + x + \cdots + x^{K-1}で割って余りが0か否かということです。

K=3として、 x^n 1+x+x^2で割って余りを求めると、
 1 \equiv 1\ mod\ (1+x+x^2)
 x \equiv x\ mod\ (1+x+x^2)
 x^2 \equiv -1-x\ mod\ (1+x+x^2)
 x^3 \equiv 1\ mod\ (1+x+x^2)
となって、K周期だということが分かります。よって、 x^nの剰余は O(K)で求められます。なので、 \displaystyle f_i = \sum_{n=0}^k{A_{n+1}x^n}\ (i = 0, \cdots , K-1)の剰余を O(NK)で作れます。これで、各クエリの多項式の剰余を O(K)で作れます。

// Match or Not
#![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()
}

fn YesNo(b: bool) -> String {
	return if b { "Yes".to_string() } else { "No".to_string() }
}


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

fn read_input() -> (usize, Vec<i64>, Vec<(usize, usize)>) {
	let v = read_vec();
	let K = v[1];
	let A: Vec<i64> = read_vec();
	let Q: usize = read();
	let queries = (0..Q).map(|_| read_vec()).map(|w| (w[0], w[1])).
									collect::<Vec<(usize, usize)>>();
	(K, A, queries)
}

// remainder of ax^n / (1+x+...+x^{K-1})
fn remainder(a: i64, n: usize, K: usize) -> Vec<i64> {
	let r = n % K;
	if r == K - 1 {
		(0..K-1).map(|_| -a).collect::<Vec<i64>>()
	}
	else {
		(0..K-1).map(|k| if k == r { a } else { 0 }).collect::<Vec<i64>>()
	}
}

fn add(f: &Vec<i64>, g: &Vec<i64>) -> Vec<i64> {
	f.iter().zip(g.iter()).map(|(&a, &b)| a + b).collect::<Vec<i64>>()
}

fn sub(f: &Vec<i64>, g: &Vec<i64>) -> Vec<i64> {
	f.iter().zip(g.iter()).map(|(&a, &b)| a - b).collect::<Vec<i64>>()
}

fn accumulate(A: &Vec<i64>, K: usize) -> Vec<Vec<i64>> {
	let mut B: Vec<Vec<i64>> = Vec::new();
	B.push((0..K).map(|_| 0).collect::<Vec<i64>>());
	for (i, a) in A.iter().enumerate() {
		let r = remainder(*a, i, K);
		let b = add(B.last().unwrap(), &r);
		B.push(b)
	}
	B
}

fn f(K: usize, A: Vec<i64>, queries: Vec<(usize, usize)>) {
	let B = accumulate(&A, K);
	for (l, r) in queries.into_iter() {
		let f = sub(&B[r], &B[l-1]);
		println!("{}", YesNo(f.iter().all(|&c| c == 0)))
	}
}

fn main() {
	let (K, A, queries) = read_input();
	f(K, A, queries)
}