https://atcoder.jp/contests/abc340/tasks/abc340_e
入力例1で、初回で2 2 0 5 6になって、2回目は6を取り出して、0に2を1～4に1を足します。このように範囲に同じ値を足すので、そういう木を作ればよいです。
// Mancala 2 #![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() } //////////////////// Segment Tree //////////////////// struct SegmentTree { n: usize, m: usize, v: Vec<i64>, } impl SegmentTree { fn add_range(&mut self, a: i64, first: usize, last: usize) { if first < last { self.add_range_core(a, (first, last), 0, 0, self.n) } } fn add_range_core(&mut self, a: i64, range: (usize, usize), i: usize, first: usize, last: usize) { let (f, l) = range; if f <= first && last <= l { self.v[i] += a } else { let mid = (first + last) / 2; if f < mid { self.add_range_core(a, range, i*2+1, first, mid) } if mid < l { self.add_range_core(a, range, i*2+2, mid, last) } } } fn get_core(&self, i: usize) -> i64 { if i == 0 { self.v[i] } else { self.v[i] + self.get_core((i-1)/2) } } fn get(&self, i: usize) -> i64 { self.get_core(i + self.n - 1) } fn set(&mut self, a: i64, i: usize) { let mut j = i + self.n - 1; let mut s = 0; loop { s += self.v[j]; if j == 0 { break } j = (j-1) / 2; } self.v[i+self.n-1] += a - s } fn print(&self) { println!("{}", (0..self.m).map(|i| self.get(i).to_string()). collect::<Vec<_>>().join(" ")) } fn ceil_two_pow(n: usize) -> usize { if n == 1 { 1 } else { SegmentTree::ceil_two_pow((n+1)/2) * 2 } } fn create(A: Vec<i64>) -> SegmentTree { let m = A.len(); let n = SegmentTree::ceil_two_pow(m); let mut v: Vec<i64> = vec![0; n*2-1]; for i in 0..m { v[i+n-1] = A[i] } SegmentTree { n, m, v } } } //////////////////// process //////////////////// fn read_input() -> (Vec<i64>, Vec<usize>) { let _v: Vec<usize> = read_vec(); let A = read_vec(); let B = read_vec(); (A, B) } fn F(A: Vec<i64>, B: Vec<usize>) { let N = A.len(); let mut tree = SegmentTree::create(A); for b in B.into_iter() { let C = tree.get(b); tree.set(0, b); let i = (b + 1) % N; let q = C / (N as i64) ; let r = (C as usize) % N; if i + r <= N { tree.add_range(q, 0, i); tree.add_range(q+1, i, i+r); tree.add_range(q, i+r, N) } else { tree.add_range(q+1, 0, i+r-N); tree.add_range(q, i+r-N, i); tree.add_range(q+1, i, N) } } tree.print() } fn main() { let (A, B) = read_input(); F(A, B) }
