AtCoder Beginner Contest 331 D

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)
}