結果
問題 | No.235 めぐるはめぐる (5) |
ユーザー | ngtkana |
提出日時 | 2024-06-24 02:41:58 |
言語 | Rust (1.77.0 + proconio) |
結果 |
WA
|
実行時間 | - |
コード長 | 24,924 bytes |
コンパイル時間 | 13,247 ms |
コンパイル使用メモリ | 403,228 KB |
実行使用メモリ | 47,680 KB |
最終ジャッジ日時 | 2024-06-24 02:42:29 |
合計ジャッジ時間 | 19,179 ms |
ジャッジサーバーID (参考情報) |
judge5 / judge3 |
(要ログイン)
テストケース
テストケース表示入力 | 結果 | 実行時間 実行使用メモリ |
---|---|---|
testcase_00 | WA | - |
testcase_01 | WA | - |
testcase_02 | WA | - |
ソースコード
use lazy_segtree::LazySegtree; use proconio::input; use proconio::marker::Usize1; fn main() { input! { n: usize, sum: [usize; n], coeff: [usize; n], edges: [(Usize1, Usize1); n - 1], q: usize, } let (hld, _g) = Hld::from_edges(0, &edges); let mut values = vec![Value::default(); n]; for (x, &i) in hld.index.iter().enumerate() { values[i] = Value { coeff: coeff[x], sum: sum[x], }; } let mut lazy_segtree = LazySegtree::<O>::from_iter(values); for _ in 0..q { input! { com: String, } match com.as_str() { "0" => { input! { i: Usize1, j: Usize1, op: usize, } hld.visit_path_segments_including_lca_by_index(i, j, |i, j| { lazy_segtree.range_apply(i..=j, &op); }); } "1" => { input! { i: Usize1, j: Usize1, } let mut ans = 0; hld.visit_path_segments_including_lca_by_index(i, j, |i, j| { ans += lazy_segtree.fold(i..=j).sum; }); println!("{ans}"); } _ => unreachable!(), } } } #[derive(Clone, Copy, Debug, PartialEq, Default)] struct Value { coeff: usize, sum: usize, } enum O {} impl lazy_segtree::Op for O { type Operator = usize; type Value = Value; fn identity() -> Self::Value { Value { coeff: 0, sum: 0 } } fn op(lhs: &Self::Value, rhs: &Self::Value) -> Self::Value { Value { coeff: lhs.coeff + rhs.coeff, sum: lhs.sum + rhs.sum, } } fn apply(op: &Self::Operator, value: &Self::Value) -> Self::Value { Value { coeff: value.coeff, sum: value.sum + op * value.coeff, } } fn identity_op() -> Self::Operator { 0 } fn compose(op: &Self::Operator, other: &Self::Operator) -> Self::Operator { op + other } } pub struct Hld { pub parent: Vec<usize>, pub index: Vec<usize>, pub head: Vec<usize>, } impl Hld { pub fn from_short_parents(mut parent: Vec<usize>) -> (Self, Vec<Vec<usize>>) { parent.insert(0, 0); let mut g = vec![Vec::new(); parent.len()]; for (i, &p) in parent.iter().enumerate().skip(1) { g[p].push(i); } (__build_hld(0, &mut g, parent), g) } pub fn from_edges(root: usize, edges: &[(usize, usize)]) -> (Self, Vec<Vec<usize>>) { let mut g = vec![Vec::new(); edges.len() + 1]; for &(i, j) in edges { g[i].push(j); g[j].push(i); } let parent = __remove_parent(root, &mut g); (__build_hld(root, &mut g, parent), g) } pub fn visit_path_segments_including_lca( &self, mut i: usize, // id mut j: usize, // id mut f: impl FnMut(usize, usize), // id ) { while self.head[i] != self.head[j] { if self.index[i] < self.index[j] { f(self.head[j], j); j = self.parent[self.head[j]]; } else { f(self.head[i], i); i = self.parent[self.head[i]]; } } if self.index[i] < self.index[j] { f(i, j) } else { f(j, i) } } pub fn visit_path_segments_including_lca_by_index( &self, i: usize, // id j: usize, // id mut f: impl FnMut(usize, usize), // index ) { self.visit_path_segments_including_lca(i, j, |i, j| f(self.index[i], self.index[j])); } pub fn lca(&self, mut i: usize, mut j: usize) -> usize { while self.head[i] != self.head[j] { if self.index[i] < self.index[j] { j = self.parent[self.head[j]]; } else { i = self.parent[self.head[i]]; } } std::cmp::min_by_key(i, j, |&i| self.index[i]) } } fn __build_hld(root: usize, g: &mut [Vec<usize>], parent: Vec<usize>) -> Hld { let n = g.len(); __heavy_first(0, g); let mut index = vec![usize::MAX; n]; let mut head = vec![usize::MAX; n]; head[root] = root; __head_and_index(0, &*g, &mut head, &mut index, &mut (0..)); Hld { parent, index, head, } } fn __head_and_index( i: usize, g: &[Vec<usize>], head: &mut [usize], index: &mut [usize], current: &mut std::ops::RangeFrom<usize>, ) { index[i] = current.next().unwrap(); for &j in &g[i] { head[j] = if j == g[i][0] { head[i] } else { j }; __head_and_index(j, g, head, index, current); } } fn __heavy_first(i: usize, g: &mut [Vec<usize>]) -> usize { let mut max = 0; let mut size = 1; for e in 0..g[i].len() { let csize = __heavy_first(g[i][e], g); if max < csize { max = csize; g[i].swap(0, e); } size += csize; } size } fn __remove_parent(root: usize, g: &mut [Vec<usize>]) -> Vec<usize> { let mut stack = vec![root]; let mut parent = vec![usize::MAX; g.len()]; parent[root] = root; while let Some(i) = stack.pop() { g[i].retain(|&j| parent[i] != j); for &j in &g[i] { parent[j] = i; stack.push(j); } } parent } // link_cut_tree {{{ // https://ngtkana.github.io/ac-adapter-rs/link_cut_tree/index.html #[allow(dead_code)] mod link_cut_tree { mod base { #[doc(hidden)] pub trait OpBase { type Value: Clone; type InternalValue: Clone; fn identity() -> Self::InternalValue; fn mul(lhs: &Self::InternalValue, rhs: &Self::InternalValue) -> Self::InternalValue; fn into_front(value: Self::InternalValue) -> Self::Value; fn from_front(value: Self::Value) -> Self::InternalValue; fn rev(value: &mut Self::InternalValue); } pub struct LinkCutTreeBase<O: OpBase> { nodes: Vec<Node<O>>, } impl<O: OpBase> LinkCutTreeBase<O> { pub fn new(n: usize) -> Self { Self { nodes: (0..n) .map(|id| Node { id, parent: std::ptr::null_mut(), left: std::ptr::null_mut(), right: std::ptr::null_mut(), rev: false, value: O::identity(), acc: O::identity(), }) .collect(), } } pub fn from_values(values: impl IntoIterator<Item = O::Value>) -> Self { Self { nodes: values .into_iter() .map(O::from_front) .enumerate() .map(|(id, value)| Node { id, parent: std::ptr::null_mut(), left: std::ptr::null_mut(), right: std::ptr::null_mut(), rev: false, value: value.clone(), acc: value, }) .collect(), } } pub fn link(&mut self, p: usize, c: usize) { unsafe { let c = std::ptr::addr_of_mut!(self.nodes[c]); let p = std::ptr::addr_of_mut!(self.nodes[p]); expose(c); assert!((*c).left.is_null(), "c = {} is not a root", (*c).id); expose(p); assert!( (*c).parent.is_null(), "c = {} and p = {} are already connected", (*c).id, (*p).id ); (*c).parent = p; (*p).right = c; update(p); } } pub fn undirected_link(&mut self, i: usize, j: usize) -> bool { if self.undirected_is_connected(i, j) { return false; } self.evert(j); self.link(i, j); true } pub fn cut(&mut self, x: usize) -> Option<usize> { unsafe { let x = std::ptr::addr_of_mut!(self.nodes[x]); expose(x); let p = (*x).left; (*x).left = std::ptr::null_mut(); let ans = p.as_ref().map(|p| p.id); if !p.is_null() { (*p).parent = std::ptr::null_mut(); } update(x); ans } } pub fn undirected_cut(&mut self, i: usize, j: usize) -> bool { if !self.undirected_has_edge(i, j) { return false; } self.evert(i); self.cut(j); true } pub fn evert(&mut self, x: usize) { unsafe { let x = std::ptr::addr_of_mut!(self.nodes[x]); expose(x); rev(x); push(x); } } pub fn undirected_has_edge(&mut self, x: usize, y: usize) -> bool { self.parent(x) == Some(y) || self.parent(y) == Some(x) } pub fn undirected_is_connected(&mut self, x: usize, y: usize) -> bool { if x == y { return true; } unsafe { let x = std::ptr::addr_of_mut!(self.nodes[x]); let y = std::ptr::addr_of_mut!(self.nodes[y]); expose(x); expose(y); !(*x).parent.is_null() } } pub fn lca(&mut self, x: usize, y: usize) -> Option<usize> { if x == y { return Some(x); } unsafe { let x = std::ptr::addr_of_mut!(self.nodes[x]); let y = std::ptr::addr_of_mut!(self.nodes[y]); expose(x); let lca = expose(y); if (*x).parent.is_null() { None } else { Some((*lca).id) } } } pub fn set(&mut self, x: usize, mut f: impl FnMut(O::Value) -> O::Value) { unsafe { let x = std::ptr::addr_of_mut!(self.nodes[x]); expose(x); (*x).value = O::from_front(f(O::into_front((*x).value.clone()))); update(x); } } pub fn fold(&mut self, x: usize) -> O::Value { unsafe { let x = std::ptr::addr_of_mut!(self.nodes[x]); expose(x); O::into_front((*x).acc.clone()) } } pub fn undirected_fold(&mut self, i: usize, j: usize) -> Option<O::Value> { if !self.undirected_is_connected(i, j) { return None; } self.evert(i); Some(self.fold(j)) } pub fn parent(&mut self, x: usize) -> Option<usize> { unsafe { let x = std::ptr::addr_of_mut!(self.nodes[x]); expose(x); let mut p = (*x).left.as_mut()?; while let Some(next) = p.right.as_mut() { p = next; } splay(p); Some(p.id) } } } #[derive(Clone, Copy)] struct Node<O: OpBase> { id: usize, parent: *mut Self, left: *mut Self, right: *mut Self, rev: bool, value: O::InternalValue, acc: O::InternalValue, } unsafe fn is_splay_root<O: OpBase>(x: *mut Node<O>) -> bool { let x = &*x; let p = match x.parent.as_ref() { Some(p) => p, None => return true, }; !std::ptr::eq(x, p.left) && !std::ptr::eq(x, p.right) } unsafe fn push<O: OpBase>(x: *mut Node<O>) { let x = &mut *x; if x.rev { if let Some(l) = x.left.as_mut() { rev(l); } if let Some(r) = x.right.as_mut() { rev(r); } x.rev = false; } } unsafe fn update<O: OpBase>(x: *mut Node<O>) { let x = &mut *x; x.acc = x.value.clone(); if !x.left.is_null() { x.acc = O::mul(&(*x.left).acc, &x.acc); } if !x.right.is_null() { x.acc = O::mul(&x.acc, &(*x.right).acc); } } unsafe fn rev<O: OpBase>(x: *mut Node<O>) { let x = &mut *x; std::mem::swap(&mut x.left, &mut x.right); O::rev(&mut x.acc); x.rev ^= true; } unsafe fn expose<O: OpBase>(x: *mut Node<O>) -> *mut Node<O> { let mut last = std::ptr::null_mut(); let mut current = x; while !current.is_null() { splay(current); (*current).right = last; update(current); last = current; current = (*current).parent; } splay(x); last } unsafe fn splay<O: OpBase>(x: *mut Node<O>) { let x = &mut *x; push(x); while !is_splay_root(x) { let p = &mut *x.parent; if is_splay_root(p) { push(p); push(x); if std::ptr::eq(p.left, x) { rotate_right(p); } else { rotate_left(p); } } else { let g = &mut *p.parent; push(g); push(p); push(x); #[allow(clippy::collapsible_else_if)] if std::ptr::eq(p.left, x) { if std::ptr::eq(g.left, p) { rotate_right(g); rotate_right(p); } else { rotate_right(p); rotate_left(g); } } else { if std::ptr::eq(g.left, p) { rotate_left(p); rotate_right(g); } else { rotate_left(g); rotate_left(p); } } } } } unsafe fn rotate_left<O: OpBase>(l: *mut Node<O>) { let l = &mut *l; let r = &mut *l.right; let p = l.parent; let c = r.left; l.right = c; if !c.is_null() { (*c).parent = l; } r.left = l; l.parent = r; r.parent = p; update(l); update(r); if !p.is_null() { if std::ptr::eq((*p).left, l) { (*p).left = r; } else if std::ptr::eq((*p).right, l) { (*p).right = r; } update(&mut *p); } } unsafe fn rotate_right<O: OpBase>(r: *mut Node<O>) { let r = &mut *r; let l = &mut *r.left; let p = r.parent; let c = l.right; r.left = c; if !c.is_null() { (*c).parent = r; } l.right = r; r.parent = l; l.parent = p; update(r); update(l); if !p.is_null() { if std::ptr::eq((*p).left, r) { (*p).left = l; } else if std::ptr::eq((*p).right, r) { (*p).right = l; } update(&mut *p); } } } pub use base::LinkCutTreeBase; use base::OpBase; pub trait Op { type Value: Clone; fn identity() -> Self::Value; fn mul(lhs: &Self::Value, rhs: &Self::Value) -> Self::Value; } impl OpBase for () { type InternalValue = (); type Value = (); fn identity() -> Self::InternalValue {} fn mul(_lhs: &Self::InternalValue, _rhs: &Self::InternalValue) -> Self::InternalValue {} fn rev(_value: &mut Self::InternalValue) {} fn into_front(_value: Self::InternalValue) {} fn from_front(_value: Self::Value) -> Self::InternalValue {} } pub type LinkCutTree = LinkCutTreeBase<()>; pub type CommutLinkCutTree<T> = LinkCutTreeBase<Commut<T>>; #[doc(hidden)] pub struct Commut<T: Op>(T); impl<T: Op> OpBase for Commut<T> { type InternalValue = T::Value; type Value = T::Value; fn identity() -> Self::InternalValue { T::identity() } fn mul(lhs: &Self::InternalValue, rhs: &Self::InternalValue) -> Self::InternalValue { T::mul(lhs, rhs) } fn rev(_value: &mut Self::InternalValue) {} fn into_front(value: Self::InternalValue) -> Self::Value { value } fn from_front(value: Self::Value) -> Self::InternalValue { value } } #[doc(hidden)] pub struct NonCommut<T: Op>(T); pub type NonCommutLinkCutTree<T> = LinkCutTreeBase<NonCommut<T>>; impl<T: Op> OpBase for NonCommut<T> { type InternalValue = (T::Value, T::Value); type Value = T::Value; fn identity() -> Self::InternalValue { (T::identity(), T::identity()) } fn mul(lhs: &Self::InternalValue, rhs: &Self::InternalValue) -> Self::InternalValue { (T::mul(&lhs.0, &rhs.0), T::mul(&rhs.1, &lhs.1)) } fn rev(value: &mut Self::InternalValue) { std::mem::swap(&mut value.0, &mut value.1); } fn into_front(value: Self::InternalValue) -> Self::Value { value.0 } fn from_front(value: Self::Value) -> Self::InternalValue { (value.clone(), value) } } } // }}} // lazy_segtree {{{ // https://ngtkana.github.io/ac-adapter-rs/lazy_segtree/index.html #[allow(dead_code)] mod lazy_segtree { use std::iter::FromIterator; use std::mem::replace; use std::ops::RangeBounds; pub trait Op { type Value; type Operator: PartialEq; fn identity() -> Self::Value; fn op(lhs: &Self::Value, rhs: &Self::Value) -> Self::Value; fn apply(op: &Self::Operator, value: &Self::Value) -> Self::Value; fn identity_op() -> Self::Operator; fn compose(op: &Self::Operator, other: &Self::Operator) -> Self::Operator; } #[derive(Debug, Clone)] pub struct LazySegtree<O: Op> { values: Vec<O::Value>, operators: Vec<O::Operator>, } impl<O: Op> LazySegtree<O> { pub fn new(values: &[O::Value]) -> Self where O::Value: Clone, O::Operator: Clone, { let values_ = values; let n = values_.len(); let mut values = vec![O::identity(); 2 * n]; values[n..].clone_from_slice(values_); for i in (1..n).rev() { values[i] = O::op(&values[i * 2], &values[i * 2 + 1]); } Self { values, operators: vec![O::identity_op(); 2 * n], } } pub fn range_apply<R: RangeBounds<usize>>(&mut self, range: R, f: &O::Operator) { let n = self.operators.len() / 2; let (l, r) = open(range, n); if l == r { return; } let l = n + l; let r = n + r; for p in (0..usize::BITS - r.leading_zeros()).rev() { self.push(l >> p); self.push((r - 1) >> p); } { let mut l = l; let mut r = r; while l < r { if l & 1 != 0 { self.operators[l] = O::compose(f, &self.operators[l]); l += 1; } if r & 1 != 0 { r -= 1; self.operators[r] = O::compose(f, &self.operators[r]); } l >>= 1; r >>= 1; } } for p in 1..usize::BITS - r.leading_zeros() { self.update(l >> p); self.update((r - 1) >> p); } } pub fn fold<R: RangeBounds<usize>>(&mut self, range: R) -> O::Value { let n = self.operators.len() / 2; let (mut l, mut r) = open(range, n); if l == r { return O::identity(); } l += n; r += n; for p in (0..usize::BITS - r.leading_zeros()).rev() { self.push(l >> p); self.push((r - 1) >> p); } for p in 1..usize::BITS - r.leading_zeros() { self.update(l >> p); self.update((r - 1) >> p); } let mut left = O::identity(); let mut right = O::identity(); while l < r { if l & 1 != 0 { left = O::op(&left, &O::apply(&self.operators[l], &self.values[l])); l += 1; } if r & 1 != 0 { r -= 1; right = O::op(&O::apply(&self.operators[r], &self.values[r]), &right); } l >>= 1; r >>= 1; } O::op(&left, &right) } pub fn get(&self, i: usize) -> O::Value where O::Value: Clone, { let mut i = self.operators.len() / 2 + i; let mut value = self.values[i].clone(); while i > 0 { value = O::apply(&self.operators[i], &value); i >>= 1; } value } pub fn collect(&self) -> Vec<O::Value> where O::Value: Clone, { (0..self.operators.len() / 2).map(|i| self.get(i)).collect() } fn push(&mut self, i: usize) { let a = replace(&mut self.operators[i], O::identity_op()); self.values[i] = O::apply(&a, &self.values[i]); if i < self.operators.len() / 2 { self.operators[i << 1] = O::compose(&a, &self.operators[i << 1]); self.operators[i << 1 | 1] = O::compose(&a, &self.operators[i << 1 | 1]); } } fn update(&mut self, i: usize) { self.values[i] = O::op( &O::apply(&self.operators[i << 1], &self.values[i << 1]), &O::apply(&self.operators[i << 1 | 1], &self.values[i << 1 | 1]), ) } } impl<O: Op> FromIterator<O::Value> for LazySegtree<O> where O::Value: Clone, O::Operator: Clone, { fn from_iter<T: IntoIterator<Item = O::Value>>(iter: T) -> Self { Self::new(&iter.into_iter().collect::<Vec<_>>()) } } fn open<B: RangeBounds<usize>>(bounds: B, n: usize) -> (usize, usize) { use std::ops::Bound; let start = match bounds.start_bound() { Bound::Unbounded => 0, Bound::Included(&x) => x, Bound::Excluded(&x) => x + 1, }; let end = match bounds.end_bound() { Bound::Unbounded => n, Bound::Included(&x) => x + 1, Bound::Excluded(&x) => x, }; (start, end) } } // }}}