https://atcoder.jp/contests/abc370/tasks/abc370_d
どの壁を爆破するかを素早く決めて、どの壁が壊れているかの情報をアップデートするために、二分木を使いたいです。縦と横でそれぞれ二分木を使います。
RustならBTreeSet<(i32, i32, bool)>で、[first, last)の範囲が壁か否かを表すために(first, last, wall)とすればよさそうですが、指定したマスがどの範囲に当たるかを調べるのが難しそうです。そのため、BTreeMap
let iter = tree.range(..x+1);
キーが[0, x+1)の範囲の要素のIteratorを取れて、その最後を取ると、xが含まれる二分木内の範囲を得られます。前後の範囲も同様に得られて、どの範囲が壁かという情報を更新できます。これを縦と横で行います。
// Gather Coins #![allow(non_snake_case)] use std::collections::BTreeMap; //////////////////// 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() } //////////////////// Point //////////////////// type Point = (i32, i32); fn read_point() -> Point { let v: Vec<i32> = read_vec(); (v[0]-1, v[1]-1) } //////////////////// Table //////////////////// type Tree = BTreeMap<i32, bool>; struct Table { rows: Vec<Tree>, cols: Vec<Tree>, } impl Table { fn explode_each(x: i32, tree: &mut Tree, L: i32) { let range = tree.range(..x+1).next_back().unwrap(); let cur_pos = *range.0; let op_next = tree.range(cur_pos+1..).next(); let next_pos = if let Some(next) = op_next { *next.0 } else { L }; if cur_pos < x && x < next_pos - 1 { // ■■■■■ -> // ■■〇■■ tree.insert(x, false); tree.insert(x + 1, true); } else if cur_pos < x && x == next_pos - 1 { // ■■■■■〇〇〇 -> // ■■■■〇〇〇〇 tree.insert(x, false); tree.remove(&(x + 1)); } else if cur_pos == x && x < next_pos - 1 { // 〇〇〇■■■■■ -> // 〇〇〇〇■■■■ tree.insert(x + 1, true); tree.remove(&x); let op_prev = tree.range(..cur_pos).next_back(); match op_prev { Some(_) => (), None => { tree.insert(x, false); } } } else if x == 0 { // ■〇〇〇 -> // 〇〇〇〇 tree.remove(&(x + 1)); tree.insert(x, false); } else { // 〇〇〇■〇〇〇 -> // 〇〇〇〇〇〇〇 tree.remove(&x); tree.remove(&(x + 1)); } } fn explode_point(&mut self, pt: Point) { let W: i32 = self.cols.len() as i32; let H: i32 = self.rows.len() as i32; let tree_row = &mut self.rows[pt.0 as usize]; Table::explode_each(pt.1, tree_row, W); let tree_col = &mut self.cols[pt.1 as usize]; Table::explode_each(pt.0, tree_col, H); } fn distances_to_walls(x: i32, cur_pos: i32, tree: &Tree) -> (i32, i32) { let d1 = if cur_pos == 0 { i32::MAX } else { x - cur_pos + 1 }; let op_next = tree.range(cur_pos+1..).next(); let d2 = match op_next { Some(rng) => rng.0 - x, None => i32::MAX }; (d1, d2) } fn get_row_range(&self, pt: Point) -> (i32, bool) { let tree = &self.rows[pt.0 as usize]; let range = tree.range(..pt.1+1).next_back().unwrap(); (*range.0, *range.1) } fn get_col_range(&self, pt: Point) -> (i32, bool) { let tree = &self.cols[pt.1 as usize]; let range = tree.range(..pt.0+1).next_back().unwrap(); (*range.0, *range.1) } fn explode(&mut self, pt: Point) { let (row_pos, wall) = self.get_row_range(pt); if wall { self.explode_point(pt) } else { let (col_pos, _) = self.get_col_range(pt); let tree_row = &mut self.rows[pt.0 as usize]; let (d1, d2) = Table::distances_to_walls(pt.1, row_pos, tree_row); let tree_col = &mut self.cols[pt.1 as usize]; let (d3, d4) = Table::distances_to_walls(pt.0, col_pos, tree_col); if d1 != i32::MAX { self.explode_point((pt.0, pt.1 - d1)) } if d2 != i32::MAX { self.explode_point((pt.0, pt.1 + d2)) } if d3 != i32::MAX { self.explode_point((pt.0 - d3, pt.1)) } if d4 != i32::MAX { self.explode_point((pt.0 + d4, pt.1)) } } } fn count(&self) -> i64 { let mut s: i64 = 0; let W: i32 = self.cols.len() as i32; for row in self.rows.iter() { let v: Vec<(i32, bool)> = row.iter(). map(|(&x, &b)| (x, b)).collect(); for i in 0..v.len()-1 { if v[i].1 { // wall s += (v[i+1].0 - v[i].0) as i64 } } if v[v.len()-1].1 { s += (W - v[v.len()-1].0) as i64 } } s } fn create(H: i32, W: i32) -> Table { let create_tree = || -> Tree { let mut tree: Tree = Tree::new(); tree.insert(0, true); tree }; let rows: Vec<Tree> = (0..H).map(|_| create_tree()).collect(); let cols: Vec<Tree> = (0..W).map(|_| create_tree()).collect(); Table { rows, cols } } } //////////////////// process //////////////////// fn read_input() -> (i32, i32, Vec<Point>) { let v: Vec<i32> = read_vec(); let (H, W, Q) = (v[0], v[1], v[2]); let points: Vec<Point> = (0..Q).map(|_| read_point()).collect(); (H, W, points) } fn F(H: i32, W: i32, points: Vec<Point>) -> i64 { let mut table = Table::create(H, W); for pt in points.into_iter() { table.explode(pt) } table.count() } fn main() { let (H, W, points) = read_input(); println!("{}", F(H, W, points)) }