Tenka1 Programmer Contest E - Equilateral

れたすさんとバチャやった時に出てきて解いた。最終的にシンプルにまとまった気がする。

問題

atcoder.jp

解法

3点の間の各マンハッタン距離が等しくなる状態をいくつか図示してみると、必ず3点のうち少なくとも2点が45度の角度で並ぶことに気づく。このままでも斜めに累積和を持つことで気合いで解けるが、45度回転すると見通しが良くなる。

コード

fn main() {
    let (r, w) = (std::io::stdin(), std::io::stdout());
    let mut sc = IO::new(r.lock(), w.lock());
    let h: usize = sc.read();
    let w: usize = sc.read();
    let board: Vec<Vec<char>> = (0..h).map(|_| sc.chars()).collect();
    let mut points = vec![];
    for i in 0..h {
        for j in 0..w {
            if board[i][j] != '#' {
                continue;
            }
            let x = i as i32;
            let y = j as i32;
            points.push((x + y, x - y));
        }
    }

    if points.is_empty() {
        println!("0");
        return;
    }

    let min_x = points.iter().map(|t| t.0).min().unwrap();
    let min_y = points.iter().map(|t| t.1).min().unwrap();

    let points = points
        .into_iter()
        .map(|(x, y)| {
            let x = (x - min_x) as usize;
            let y = (y - min_y) as usize;
            (x, y)
        })
        .collect::<Vec<_>>();

    let max_x = points.iter().map(|t| t.0).max().unwrap();
    let max_y = points.iter().map(|t| t.1).max().unwrap();

    let h = max_x + 1;
    let w = max_y + 1;
    let mut board = vec![vec![false; w]; h];
    for &(x, y) in points.iter() {
        board[x][y] = true;
    }

    let mut row_sum = vec![vec![0; w + 1]; h];
    for i in 0..h {
        for j in 0..w {
            row_sum[i][j + 1] = row_sum[i][j] + if board[i][j] { 1 } else { 0 };
        }
    }

    let mut col_sum = vec![vec![0; h + 1]; w];
    for j in 0..w {
        for i in 0..h {
            col_sum[j][i + 1] = col_sum[j][i] + if board[i][j] { 1 } else { 0 };
        }
    }

    let mut ans = 0;
    for i in 0..h {
        for j2 in 0..w {
            if !board[i][j2] {
                continue;
            }
            for j1 in 0..j2 {
                if !board[i][j1] {
                    continue;
                }

                let d = j2 - j1;
                if i + d < h {
                    let add = row_sum[i + d][j2] - row_sum[i + d][j1 + 1];
                    ans += add;
                }
                if i >= d {
                    let add = row_sum[i - d][j2] - row_sum[i - d][j1 + 1];
                    ans += add;
                }
            }
        }
    }

    for j in 0..w {
        for i2 in 0..h {
            if !board[i2][j] {
                continue;
            }
            for i1 in 0..i2 {
                if !board[i1][j] {
                    continue;
                }

                let d = i2 - i1;
                if j + d < w {
                    let add = col_sum[j + d][i2] - col_sum[j + d][i1 + 1];
                    ans += add;
                }
                if j >= d {
                    let add = col_sum[j - d][i2] - col_sum[j - d][i1 + 1];
                    ans += add;
                }
            }
        }
    }

    for i in 0..h {
        for j2 in 0..w {
            if !board[i][j2] {
                continue;
            }
            for j1 in 0..j2 {
                if !board[i][j1] {
                    continue;
                }

                let d = j2 - j1;
                if i + d < h && board[i + d][j1] {
                    ans += 1;
                }
                if i + d < h && board[i + d][j2] {
                    ans += 1;
                }
                if i >= d && board[i - d][j1] {
                    ans += 1;
                }
                if i >= d && board[i - d][j2] {
                    ans += 1;
                }
            }
        }
    }

    println!("{}", ans);
}

pub struct IO<R, W: std::io::Write>(R, std::io::BufWriter<W>);

impl<R: std::io::Read, W: std::io::Write> IO<R, W> {
    pub fn new(r: R, w: W) -> Self {
        Self(r, std::io::BufWriter::new(w))
    }
    pub fn write<S: ToString>(&mut self, s: S) {
        use std::io::Write;
        self.1.write_all(s.to_string().as_bytes()).unwrap();
    }
    pub fn read<T: std::str::FromStr>(&mut self) -> T {
        use std::io::Read;
        let buf = self
            .0
            .by_ref()
            .bytes()
            .map(|b| b.unwrap())
            .skip_while(|&b| b == b' ' || b == b'\n' || b == b'\r' || b == b'\t')
            .take_while(|&b| b != b' ' && b != b'\n' && b != b'\r' && b != b'\t')
            .collect::<Vec<_>>();
        unsafe { std::str::from_utf8_unchecked(&buf) }
            .parse()
            .ok()
            .expect("Parse error.")
    }
    pub fn vec<T: std::str::FromStr>(&mut self, n: usize) -> Vec<T> {
        (0..n).map(|_| self.read()).collect()
    }
    pub fn chars(&mut self) -> Vec<char> {
        self.read::<String>().chars().collect()
    }
}