https://atcoder.jp/contests/abc331/tasks/abc331_d
二次元累積和を使うと高速にエリア内の黒の個数を求めることができます。
// Tile Pattern #![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() } //////////////////// Rect //////////////////// type Rect = (usize, usize, usize, usize); //////////////////// Two-Dimensional-Cumulative-Sum //////////////////// fn make_TDCS(table: Vec<Vec<i64>>) -> Vec<Vec<i64>> { let H = table.len(); let W = table[0].len(); let mut TDCS: Vec<Vec<i64>> = vec![vec![0; W+1]; H+1]; for i in 0..H { for j in 0..W { TDCS[i+1][j+1] = TDCS[i+1][j] + TDCS[i][j+1] - TDCS[i][j] + table[i][j] } } TDCS } fn sum_by_area(rect: Rect, TDCS: &Vec<Vec<i64>>) -> i64 { TDCS[rect.2][rect.3] - TDCS[rect.2][rect.1] - TDCS[rect.0][rect.3] + TDCS[rect.0][rect.1] } //////////////////// process //////////////////// fn read_input() -> (usize, Vec<String>) { let v = read_vec(); let (N, Q) = (v[0], v[1]); let P: Vec<String> = (0..N).map(|_| read::<String>()).collect(); (Q, P) } fn read_query() -> Rect { let v = read_vec(); let (A, B, C, D) = (v[0], v[1], v[2]+1, v[3]+1); (A, B, C, D) } fn count_blacks(rect: Rect, TDCS: &Vec<Vec<i64>>) -> i64 { let N = TDCS.len() - 1; let dx1 = (rect.0 + N - 1) / N * N - rect.0; let x2 = rect.2 % N; let dy1 = (rect.1 + N - 1) / N * N - rect.1; let y2 = rect.3 % N; let x1 = if dx1 == 0 { 0 } else { N - dx1 }; let y1 = if dy1 == 0 { 0 } else { N - dy1 }; if rect.0 / N == rect.2 / N && rect.1 / N == rect.3 / N { sum_by_area((x1, y1, x2, y2), TDCS) } else if rect.0 / N == rect.2 / N { let ny = ((rect.3 - rect.1 - dy1 - y2) / N) as i64; sum_by_area((x1, N-dy1, x2, N), TDCS) + sum_by_area((x1, 0, x2, N), TDCS) * ny + sum_by_area((x1, 0, x2, y2), TDCS) } else if rect.1 / N == rect.3 / N { let nx = ((rect.2 - rect.0 - dx1 - x2) / N) as i64; sum_by_area((N-dx1, y1, N, y2), TDCS) + sum_by_area((0, y1, N, y2), TDCS) * nx + sum_by_area((0, y1, x2, y2), TDCS) } else { let nx = ((rect.2 - rect.0 - dx1 - x2) / N) as i64; let ny = ((rect.3 - rect.1 - dy1 - y2) / N) as i64; let counter1 = sum_by_area((0, 0, N, N), TDCS) * nx * ny; let counter2 = sum_by_area((N-dx1, N-dy1, N, N), TDCS); let counter3 = sum_by_area((N-dx1, 0, N, y2), TDCS); let counter4 = sum_by_area((0, N-dy1, x2, N), TDCS); let counter5 = sum_by_area((0, 0, x2, y2), TDCS); let counter6 = sum_by_area((0, N-dy1, N, N), TDCS) * nx; let counter7 = sum_by_area((N-dx1, 0, N, N), TDCS) * ny; let counter8 = sum_by_area((0, 0, N, y2), TDCS) * nx; let counter9 = sum_by_area((0, 0, x2, N), TDCS) * ny; counter1 + counter2 + counter3 + counter4 + counter5 + counter6 + counter7 + counter8 + counter9 } } fn F(Q: usize, P: Vec<String>) { let table: Vec<Vec<i64>> = P.into_iter(). map(|s| s.chars(). map(|c| if c == 'B' { 1 } else { 0 }). collect::<Vec<i64>>()). collect(); let TDCS = make_TDCS(table); for _ in 0..Q { let rect = read_query(); println!("{}", count_blacks(rect, &TDCS)) } } fn main() { let (Q, P) = read_input(); F(Q, P) }