AtCoder Beginner Contest 366 D

https://atcoder.jp/contests/abc366/tasks/abc366_d

を計算しておけば包除原理で求められますが、このPを求めるときも包除原理を使います。

// Cuboid Sum Query
#![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()
}


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

fn read_input() -> (Vec<Vec<Vec<i32>>>, usize) {
    let N: usize = read();
    let A: Vec<Vec<Vec<i32>>> = (0..N).map(|_|
                                (0..N).map(|_| read_vec::<i32>()).
                                                collect::<Vec<Vec<i32>>>()).
                                                collect();
    let Q: usize = read();
    (A, Q)
}

fn accumulate(A: Vec<Vec<Vec<i32>>>) -> Vec<Vec<Vec<i32>>> {
    let N = A.len();
    let mut P: Vec<Vec<Vec<i32>>> = vec![vec![vec![0; N+1]; N+1]; N+1];
    P[0][0][0] = A[0][0][0];
    for x in 0..N {
        for y in 0..N {
            for z in 0..N {
                P[x+1][y+1][z+1] = A[x][y][z]
                                    + P[x][y+1][z+1]
                                    + P[x+1][y][z+1]
                                    + P[x+1][y+1][z]
                                    - P[x][y][z+1]
                                    - P[x][y+1][z]
                                    - P[x+1][y][z]
                                    + P[x][y][z];
            }
        }
    }
    P
}

fn F_each(v: Vec<usize>, P: &Vec<Vec<Vec<i32>>>) -> i32 {
    let (Lx, Rx, Ly, Ry, Lz, Rz) = (v[0], v[1], v[2], v[3], v[4], v[5]);
    P[Rx][Ry][Rz] - P[Lx-1][Ry][Rz] - P[Rx][Ly-1][Rz] - P[Rx][Ry][Lz-1] +
                    P[Lx-1][Ly-1][Rz] + P[Lx-1][Ry][Lz-1] + P[Rx][Ly-1][Lz-1] -
                    P[Lx-1][Ly-1][Lz-1]
}

fn F(A: Vec<Vec<Vec<i32>>>, Q: usize) {
    let P = accumulate(A);
    for _ in 0..Q {
        let LR: Vec<usize> = read_vec();
        println!("{}", F_each(LR, &P))
    }
}

fn main() {
    let (A, Q) = read_input();
    F(A, Q)
}