結果

問題 No.900 aδδitivee
ユーザー ngtkana
提出日時 2024-05-07 18:41:49
言語 Rust
(1.83.0 + proconio)
結果
AC  
実行時間 154 ms / 2,000 ms
コード長 21,846 bytes
コンパイル時間 16,561 ms
コンパイル使用メモリ 382,424 KB
実行使用メモリ 24,704 KB
最終ジャッジ日時 2024-11-30 13:40:42
合計ジャッジ時間 23,006 ms
ジャッジサーバーID
(参考情報)
judge5 / judge1
このコードへのチャレンジ
(要ログイン)
ファイルパターン 結果
sample AC * 2
other AC * 27
権限があれば一括ダウンロードができます

ソースコード

diff #
プレゼンテーションモードにする

use crate::segtree::Segtree;
use proconio::input;
use std::cmp::Reverse;
use std::ops::Bound;
fn main() {
input! {
n: usize,
edges: [(usize, usize, i64); n - 1],
q: usize,
}
let mut g = vec![Vec::new(); n];
for &(i, j, _) in &edges {
g[i].push(j);
g[j].push(i);
}
let mut stack = vec![0];
let mut parent = vec![usize::MAX; n];
let mut depth = vec![0; n];
parent[0] = 0;
let mut size = vec![1; n];
while let Some(x) = stack.pop() {
if x <= isize::MAX as usize {
stack.push(!x);
g[x].retain(|&y| y != parent[x]);
for &y in &g[x] {
if parent[y] != usize::MAX {
continue;
}
parent[y] = x;
depth[y] = depth[x] + 1;
stack.push(y);
}
} else {
let x = !x;
g[x].sort_unstable_by_key(|&y| Reverse(size[y]));
for &y in &g[x] {
size[x] += size[y];
}
}
}
let mut sorted = Vec::new();
let mut position = vec![usize::MAX; n];
let mut head = (0..n).collect::<Vec<_>>();
let mut stack = vec![0];
while let Some(x) = stack.pop() {
position[x] = sorted.len();
sorted.push(x);
if let Some(&y) = g[x].first() {
head[y] = head[x];
}
stack.extend(g[x].iter().rev().copied());
}
let mut original = Segtree::<O1>::from_len(n);
let mut additional = Segtree::<O2>::from_len(n);
for &(i, j, w) in &edges {
let i = if parent[i] == j { i } else { j };
*original.entry(position[i]) = w;
}
for _ in 0..q {
input! {
com: u8,
}
match com {
1 => {
input! {
i: usize,
x: i64,
}
let mut e = additional.entry(position[i]);
e.additional += x;
e.weighted += x * depth[i] as i64;
}
2 => {
input! {
mut j: usize,
}
let mut ans = 0;
let coeff = depth[j];
loop {
if head[j] == 0 {
let original =
original.fold((Bound::Excluded(0), Bound::Included(position[j])));
let Value {
additional,
weighted,
} = additional.fold(0..=position[j]);
ans += additional * coeff as i64 - weighted + original;
break;
}
let original = original.fold(position[head[j]]..=position[j]);
let Value {
additional,
weighted,
} = additional.fold(position[head[j]]..=position[j]);
ans += additional * coeff as i64 - weighted + original;
j = parent[head[j]];
}
println!("{}", ans);
}
_ => unreachable!(),
}
}
}
enum O1 {}
impl segtree::Op for O1 {
type Value = i64;
fn identity() -> Self::Value {
0
}
fn op(lhs: &Self::Value, rhs: &Self::Value) -> Self::Value {
lhs + rhs
}
}
#[derive(Clone, Copy)]
struct Value {
additional: i64,
weighted: i64,
}
enum O2 {}
impl segtree::Op for O2 {
type Value = Value;
fn identity() -> Self::Value {
Value {
additional: 0,
weighted: 0,
}
}
fn op(lhs: &Self::Value, rhs: &Self::Value) -> Self::Value {
Value {
additional: lhs.additional + rhs.additional,
weighted: lhs.weighted + rhs.weighted,
}
}
}
// segtree {{{
// https://ngtkana.github.io/ac-adapter-rs/segtree/index.html
#[allow(dead_code)]
mod segtree {
use core::fmt;
use std::collections::BTreeMap;
use std::iter::FromIterator;
use std::ops::Deref;
use std::ops::DerefMut;
use std::ops::Index;
use std::ops::RangeBounds;
pub trait Op {
type Value;
fn identity() -> Self::Value;
fn op(lhs: &Self::Value, rhs: &Self::Value) -> Self::Value;
}
pub struct Segtree<O: Op> {
values: Vec<O::Value>,
}
impl<O: Op> Segtree<O> {
pub fn from_len(n: usize) -> Self
where
O::Value: Clone,
{
Self::new(&vec![O::identity(); n])
}
pub fn new(values: &[O::Value]) -> Self
where
O::Value: 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 }
}
pub fn fold<R: RangeBounds<usize>>(&self, range: R) -> O::Value {
let n = self.values.len() / 2;
let (mut start, mut end) = open(range, n);
assert!(start <= end && end <= n);
start += n;
end += n;
let mut left = O::identity();
let mut right = O::identity();
while start < end {
if start % 2 == 1 {
left = O::op(&left, &self.values[start]);
start += 1;
}
if end % 2 == 1 {
end -= 1;
right = O::op(&self.values[end], &right);
}
start /= 2;
end /= 2;
}
O::op(&left, &right)
}
pub fn entry(&mut self, index: usize) -> Entry<O> {
let n = self.values.len() / 2;
Entry {
segtree: self,
index: n + index,
}
}
pub fn iter(&self) -> impl Iterator<Item = &O::Value> {
self.values[self.values.len() / 2..].iter()
}
pub fn as_slice(&self) -> &[O::Value] {
&self.values[self.values.len() / 2..]
}
}
impl<O: Op> fmt::Debug for Segtree<O>
where
O::Value: fmt::Debug,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Segtree")
.field("values", &self.values)
.finish()
}
}
impl<O: Op> FromIterator<O::Value> for Segtree<O>
where
O::Value: Clone,
{
fn from_iter<I: IntoIterator<Item = O::Value>>(iter: I) -> Self {
Self::new(&iter.into_iter().collect::<Vec<_>>())
}
}
impl<O: Op> Index<usize> for Segtree<O> {
type Output = O::Value;
fn index(&self, index: usize) -> &Self::Output {
&self.values[self.values.len() / 2 + index]
}
}
pub struct Entry<'a, O: Op> {
segtree: &'a mut Segtree<O>,
index: usize,
}
impl<'a, O: Op> Drop for Entry<'a, O> {
fn drop(&mut self) {
let mut index = self.index;
while index != 0 {
index /= 2;
self.segtree.values[index] = O::op(
&self.segtree.values[index * 2],
&self.segtree.values[index * 2 + 1],
);
}
}
}
impl<'a, O: Op> Deref for Entry<'a, O> {
type Target = O::Value;
fn deref(&self) -> &Self::Target {
&self.segtree.values[self.index]
}
}
impl<'a, O: Op> DerefMut for Entry<'a, O> {
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.segtree.values[self.index]
}
}
impl<'a, O: Op> fmt::Debug for Entry<'a, O>
where
O::Value: fmt::Debug,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Entry").field("index", &self.index).finish()
}
}
pub struct SparseSegtree<K, O: Op> {
inner: Segtree<O>,
keys: Vec<K>,
}
impl<K: Ord, O: Op> SparseSegtree<K, O> {
pub fn new(kv: &[(K, O::Value)]) -> Self
where
K: Clone,
O::Value: Clone,
{
let mut keys = kv.iter().map(|(k, _)| k.clone()).collect::<Vec<_>>();
keys.sort();
let values = kv.iter().map(|(_, v)| v.clone()).collect::<Vec<_>>();
Self {
inner: Segtree::new(&values),
keys,
}
}
pub fn fold<R: RangeBounds<K>>(&self, range: R) -> O::Value {
let (start, end) = open_key(range, &self.keys);
self.inner.fold(start..end)
}
pub fn entry(&mut self, key: &K) -> Entry<'_, O> {
let index = self.keys.binary_search(key).unwrap() + self.keys.len();
Entry {
segtree: &mut self.inner,
index,
}
}
pub fn keys(&self) -> &[K] {
&self.keys
}
pub fn iter(&self) -> impl Iterator<Item = (&K, &O::Value)> {
self.keys.iter().zip(self.inner.as_slice())
}
pub fn collect_map(&self) -> BTreeMap<K, O::Value>
where
K: Clone,
O::Value: Clone,
{
self.keys
.iter()
.cloned()
.zip(self.inner.iter().cloned())
.collect()
}
}
impl<K, O: Op> fmt::Debug for SparseSegtree<K, O>
where
K: fmt::Debug,
O::Value: fmt::Debug,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("SparseSegtree")
.field("inner", &self.inner)
.field("keys", &self.keys)
.finish()
}
}
impl<K: Ord, O: Op> FromIterator<(K, O::Value)> for SparseSegtree<K, O>
where
K: Clone,
O::Value: Clone,
{
fn from_iter<I: IntoIterator<Item = (K, O::Value)>>(iter: I) -> Self {
Self::new(&iter.into_iter().collect::<Vec<_>>())
}
}
impl<K: Ord, O: Op> Index<K> for SparseSegtree<K, O> {
type Output = O::Value;
fn index(&self, key: K) -> &Self::Output {
&self.inner[self.keys.binary_search(&key).unwrap()]
}
}
pub struct Sparse2dSegtree<K, L, O: Op> {
segtrees: Vec<SparseSegtree<L, O>>,
keys: Vec<K>,
}
impl<K, L, O: Op> Sparse2dSegtree<K, L, O>
where
K: Ord + Clone,
L: Ord + Clone,
O::Value: Clone,
{
pub fn new(points: &[(K, L, O::Value)]) -> Self {
let mut keys = points.iter().map(|(k, _, _)| k.clone()).collect::<Vec<_>>();
keys.sort();
keys.dedup();
let mut lvs = vec![vec![]; keys.len() * 2];
for (k, l, v) in points {
let mut i = keys.binary_search(k).unwrap();
i += keys.len();
while i != 0 {
lvs[i].push((l.clone(), v.clone()));
i /= 2;
}
}
let segtrees = lvs
.into_iter()
.map(|lvs_| {
let mut ls = lvs_.iter().map(|(l, _)| l).collect::<Vec<_>>();
ls.sort();
ls.dedup();
let mut lvs = ls
.iter()
.map(|&l| (l.clone(), O::identity()))
.collect::<Vec<_>>();
for (l, v) in &lvs_ {
let i = ls.binary_search(&l).unwrap();
lvs[i].1 = O::op(&lvs[i].1, v);
}
SparseSegtree::new(&lvs)
})
.collect::<Vec<_>>();
Self { segtrees, keys }
}
pub fn fold(&self, i: impl RangeBounds<K>, j: impl RangeBounds<L> + Clone) -> O::Value {
let (mut i0, mut i1) = open_key(i, &self.keys);
i0 += self.keys.len();
i1 += self.keys.len();
let mut left = O::identity();
let mut right = O::identity();
while i0 < i1 {
if i0 % 2 == 1 {
left = O::op(&left, &self.segtrees[i0].fold(j.clone()));
i0 += 1;
}
if i1 % 2 == 1 {
i1 -= 1;
right = O::op(&self.segtrees[i1].fold(j.clone()), &right);
}
i0 /= 2;
i1 /= 2;
}
O::op(&left, &right)
}
pub fn apply(&mut self, k: &K, l: &L, mut f: impl FnMut(&mut O::Value)) {
let mut i = self.keys.binary_search(k).unwrap();
i += self.keys.len();
while i != 0 {
f(&mut self.segtrees[i].entry(l));
i /= 2;
}
}
pub fn iter(&self) -> impl Iterator<Item = (&K, &L, &O::Value)> {
self.keys
.iter()
.zip(self.segtrees[self.keys.len()..].iter())
.flat_map(|(k, segtree)| segtree.iter().map(move |(l, v)| (k, l, v)))
}
pub fn collect_map(&self) -> BTreeMap<(K, L), O::Value>
where
K: Clone,
L: Clone,
O::Value: Clone,
{
self.keys
.iter()
.flat_map(|k| {
self.segtrees[self.keys.len() + self.keys.binary_search(k).unwrap()]
.iter()
.map(move |(l, v)| ((k.clone(), l.clone()), v.clone()))
})
.collect()
}
}
impl<K, L, O: Op> fmt::Debug for Sparse2dSegtree<K, L, O>
where
K: fmt::Debug,
L: fmt::Debug,
O::Value: fmt::Debug,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Sparse2dSegtree")
.field("segtrees", &self.segtrees)
.field("keys", &self.keys)
.finish()
}
}
impl<K, L, O: Op> FromIterator<(K, L, O::Value)> for Sparse2dSegtree<K, L, O>
where
K: Ord + Clone,
L: Ord + Clone,
O::Value: Clone,
{
fn from_iter<I: IntoIterator<Item = (K, L, O::Value)>>(iter: I) -> Self {
Self::new(&iter.into_iter().collect::<Vec<_>>())
}
}
impl<K: Ord, L: Ord, O: Op> Index<K> for Sparse2dSegtree<K, L, O> {
type Output = SparseSegtree<L, O>;
fn index(&self, i: K) -> &Self::Output {
&self.segtrees[self.keys.binary_search(&i).unwrap() + self.keys.len()]
}
}
impl<K: Ord, L: Ord, O: Op> Index<(K, L)> for Sparse2dSegtree<K, L, O> {
type Output = O::Value;
fn index(&self, (i, j): (K, L)) -> &Self::Output {
&self.segtrees[self.keys.binary_search(&i).unwrap() + self.keys.len()][j]
}
}
pub struct Dense2dSegtree<O: Op> {
values: Vec<Vec<O::Value>>,
}
impl<O: Op> Dense2dSegtree<O> {
pub fn new(values: &[Vec<O::Value>]) -> Self
where
O::Value: Clone,
{
let values_ = values;
let h = values.len();
let w = values.get(0).map_or(0, Vec::len);
let mut values = vec![vec![O::identity(); 2 * w]; 2 * h];
for (values, values_) in values[h..].iter_mut().zip(values_) {
values[w..].clone_from_slice(values_);
for j in (1..w).rev() {
values[j] = O::op(&values[j * 2], &values[j * 2 + 1]);
}
}
for i in (1..h).rev() {
for j in 0..2 * w {
values[i][j] = O::op(&values[i * 2][j], &values[i * 2 + 1][j]);
}
}
Self { values }
}
pub fn fold(&self, i: impl RangeBounds<usize>, j: impl RangeBounds<usize>) -> O::Value {
let h = self.values.len() / 2;
let w = self.values.get(0).map_or(0, |v| v.len() / 2);
let (mut i0, mut i1) = open(i, h);
assert!(i0 <= i1 && i1 <= h);
let (mut j0, mut j1) = open(j, w);
assert!(j0 <= j1 && j1 <= w);
i0 += h;
i1 += h;
j0 += w;
j1 += w;
let mut left = O::identity();
let mut right = O::identity();
while i0 < i1 {
if i0 % 2 == 1 {
let mut j0 = j0;
let mut j1 = j1;
while j0 < j1 {
if j0 % 2 == 1 {
left = O::op(&left, &self.values[i0][j0]);
j0 += 1;
}
if j1 % 2 == 1 {
j1 -= 1;
right = O::op(&self.values[i0][j1], &right);
}
j0 /= 2;
j1 /= 2;
}
i0 += 1;
}
if i1 % 2 == 1 {
i1 -= 1;
let mut j0 = j0;
let mut j1 = j1;
while j0 < j1 {
if j0 % 2 == 1 {
left = O::op(&left, &self.values[i1][j0]);
j0 += 1;
}
if j1 % 2 == 1 {
j1 -= 1;
right = O::op(&self.values[i1][j1], &right);
}
j0 /= 2;
j1 /= 2;
}
}
i0 /= 2;
i1 /= 2;
}
O::op(&left, &right)
}
pub fn entry(&mut self, i: usize, j: usize) -> Dense2dEntry<O> {
let h = self.values.len() / 2;
let w = self.values.get(0).map_or(0, |v| v.len() / 2);
Dense2dEntry {
segtree: self,
i: h + i,
j: w + j,
}
}
pub fn iter(&self) -> impl Iterator<Item = &[O::Value]> {
self.values[self.values.len() / 2..]
.iter()
.map(|v| &v[v.len() / 2..])
}
pub fn collect_vec(&self) -> Vec<Vec<O::Value>>
where
O::Value: Clone,
{
self.iter().map(<[_]>::to_vec).collect()
}
}
impl<O: Op> fmt::Debug for Dense2dSegtree<O>
where
O::Value: fmt::Debug,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Dense2dSegtree")
.field("values", &self.values)
.finish()
}
}
impl<O: Op> Index<usize> for Dense2dSegtree<O> {
type Output = [O::Value];
fn index(&self, index: usize) -> &Self::Output {
&self.values[self.values.len() / 2 + index]
}
}
pub struct Dense2dEntry<'a, O: Op> {
segtree: &'a mut Dense2dSegtree<O>,
i: usize,
j: usize,
}
impl<'a, O: Op> Drop for Dense2dEntry<'a, O> {
fn drop(&mut self) {
let mut i = self.i;
let mut j = self.j / 2;
while j != 0 {
self.segtree.values[i][j] = O::op(
&self.segtree.values[i][2 * j],
&self.segtree.values[i][2 * j + 1],
);
j /= 2;
}
i /= 2;
while i != 0 {
let mut j = self.j;
while j != 0 {
self.segtree.values[i][j] = O::op(
&self.segtree.values[i * 2][j],
&self.segtree.values[i * 2 + 1][j],
);
j /= 2;
}
i /= 2;
}
}
}
impl<'a, O: Op> Deref for Dense2dEntry<'a, O> {
type Target = O::Value;
fn deref(&self) -> &Self::Target {
&self.segtree.values[self.i][self.j]
}
}
impl<'a, O: Op> DerefMut for Dense2dEntry<'a, O> {
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.segtree.values[self.i][self.j]
}
}
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)
}
fn open_key<K: Ord, B: RangeBounds<K>>(bounds: B, keys: &[K]) -> (usize, usize) {
use std::ops::Bound;
let start = match bounds.start_bound() {
Bound::Unbounded => 0,
Bound::Included(x) => match keys.binary_search(x) {
Ok(i) | Err(i) => i,
},
Bound::Excluded(x) => match keys.binary_search(x) {
Ok(i) => i + 1,
Err(i) => i,
},
};
let end = match bounds.end_bound() {
Bound::Unbounded => keys.len(),
Bound::Included(x) => match keys.binary_search(x) {
Ok(i) => i + 1,
Err(i) => i,
},
Bound::Excluded(x) => match keys.binary_search(x) {
Ok(i) | Err(i) => i,
},
};
(start, end)
}
}
// }}}
הההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההה
XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
0