結果
| 問題 |
No.899 γatheree
|
| コンテスト | |
| ユーザー |
 Haar
|
| 提出日時 | 2022-09-06 16:54:59 |
| 言語 | Rust (1.83.0 + proconio) |
| 結果 |
AC
|
| 実行時間 | 1,595 ms / 2,000 ms |
| コード長 | 19,113 bytes |
| コンパイル時間 | 13,642 ms |
| コンパイル使用メモリ | 383,928 KB |
| 実行使用メモリ | 29,676 KB |
| 最終ジャッジ日時 | 2024-11-21 18:01:11 |
| 合計ジャッジ時間 | 40,149 ms |
|
ジャッジサーバーID （参考情報） |
judge4 / judge1 |
（要ログイン）
| ファイルパターン | 結果 |
|---|---|
| sample | AC * 1 |
| other | AC * 23 |
ソースコード
// Bundled at 2022/09/06 16:54:35 +09:00
// Author: Haar
pub mod main {
use super::*;
use haar_lib::{
algebra::update_sum::*,
ds::lazy_segtree::*,
tree::{depth_query::*, *},
utils::fastio::*,
};
#[derive(Clone, Default)]
pub struct Problem {}
impl Problem {
pub fn main(&mut self) -> Result<(), Box<dyn std::error::Error>> {
let mut io = FastIO::new();
let n = io.read_u64() as usize;
let mut tree = Tree::new(n);
for _ in 0..n - 1 {
let u = io.read_u64() as usize;
let v = io.read_u64() as usize;
tree.extend(Some(TreeEdge::new(u, v, (), ())));
}
let res = TreeDepthQuery::new(&tree, 0);
let a = (0..n).map(|_| io.read_u64()).collect::<Vec<_>>();
let mut seg = LazySegmentTree::new(n, UpdateSum::<u64, u64>::new());
for i in 0..n {
let (l, r) = res.me_query(i);
seg.update(l..r, Some(a[i]));
}
let q = io.read_u64() as usize;
for _ in 0..q {
let x = io.read_u64() as usize;
let mut ans = 0;
let mut f = |(l, r)| {
ans += seg.fold(l..r);
seg.update(l..r, Some(0));
};
if let Some(x) = res.ancestor(x, 2) {
f(res.me_query(x));
}
if let Some(x) = res.ancestor(x, 1) {
f(res.me_query(x));
res.children_query(x, 1).into_iter().for_each(|p| f(p));
}
f(res.me_query(x));
res.children_query(x, 1).into_iter().for_each(|p| f(p));
res.children_query(x, 2).into_iter().for_each(|p| f(p));
let (l, r) = res.me_query(x);
seg.update(l..r, Some(ans));
io.writeln(ans);
}
Ok(())
}
}
}
fn main() {
main::Problem::default().main().unwrap();
}
use crate as haar_lib;
pub mod algebra {
pub mod action {
pub trait Action {
type FType;
type UType;
fn fold_id(&self) -> Self::FType;
fn fold(&self, x: Self::FType, y: Self::FType) -> Self::FType;
fn update_id(&self) -> Self::UType;
fn update(&self, next: Self::UType, cur: Self::UType) -> Self::UType;
fn convert(&self, x: Self::FType, y: Self::UType, l: usize) -> Self::FType;
}
}
pub mod update_sum {
use crate::algebra::action::Action;
use std::{
marker::PhantomData,
ops::{Add, Mul},
};
#[derive(Clone, Default)]
pub struct UpdateSum<T, U>(PhantomData<T>, PhantomData<U>);
impl<T, U> UpdateSum<T, U> {
pub fn new() -> Self {
Self(PhantomData, PhantomData)
}
}
impl<T, U> Action for UpdateSum<T, U>
where
T: Add<Output = T> + Default + From<U>,
U: Mul<Output = U> + Default + From<u64>,
{
type FType = T;
type UType = Option<U>;
fn fold_id(&self) -> Self::FType {
T::default()
}
fn fold(&self, x: Self::FType, y: Self::FType) -> Self::FType {
x + y
}
fn update_id(&self) -> Self::UType {
None
}
fn update(&self, x: Self::UType, y: Self::UType) -> Self::UType {
match x {
Some(_) => x,
_ => y,
}
}
fn convert(&self, x: Self::FType, y: Self::UType, l: usize) -> Self::FType {
match y {
Some(y) => T::from(y * U::from(l as u64)),
_ => x,
}
}
}
}
}
pub mod ds {
pub mod traits {
pub trait Foldable<Idx> {
type Output;
fn fold(&self, range: Idx) -> Self::Output;
}
pub trait Foldable2D<Idx> {
type Output;
fn fold(&self, x_range: Idx, y_range: Idx) -> Self::Output;
}
pub trait Assignable<Idx> {
type Value;
fn assign(&mut self, i: Idx, value: Self::Value);
}
pub trait Updatable<Idx> {
type Value;
fn update(&mut self, i: Idx, value: Self::Value);
}
pub trait Indexable<Idx> {
type Output;
fn get(&self, i: Idx) -> Self::Output;
}
}
pub mod lazy_segtree {
use crate::algebra::action::Action;
pub use crate::ds::traits::Updatable;
use std::ops::Range;
pub struct LazySegmentTree<T, U, A> {
size: usize,
data: Vec<T>,
lazy: Vec<U>,
action: A,
}
impl<T: Clone + Eq, U: Clone + Eq, A: Clone + Action<FType = T, UType = U>>
LazySegmentTree<T, U, A>
{
pub fn new(n: usize, a: A) -> Self {
let size = n.next_power_of_two() * 2;
Self {
size,
data: vec![a.fold_id(); size],
lazy: vec![a.update_id(); size],
action: a,
}
}
fn propagate(&mut self, i: usize) {
if self.lazy[i] == self.action.update_id() {
return;
}
if i < self.size / 2 {
self.lazy[i << 1] = self
.action
.update(self.lazy[i].clone(), self.lazy[i << 1].clone());
self.lazy[i << 1 | 1] = self
.action
.update(self.lazy[i].clone(), self.lazy[i << 1 | 1].clone());
}
let len = (self.size / 2) >> (31 - (i as u32).leading_zeros());
self.data[i] = self
.action
.convert(self.data[i].clone(), self.lazy[i].clone(), len);
self.lazy[i] = self.action.update_id();
}
fn propagate_top_down(&mut self, mut i: usize) {
let mut temp = vec![];
while i > 1 {
i >>= 1;
temp.push(i);
}
for &i in temp.iter().rev() {
self.propagate(i);
}
}
fn bottom_up(&mut self, mut i: usize) {
while i > 1 {
i >>= 1;
self.propagate(i << 1);
self.propagate(i << 1 | 1);
self.data[i] = self
.action
.fold(self.data[i << 1].clone(), self.data[i << 1 | 1].clone());
}
}
pub fn fold(&mut self, Range { start: l, end: r }: Range<usize>) -> T {
self.propagate_top_down(l + self.size / 2);
if r < self.size / 2 {
self.propagate_top_down(r + self.size / 2);
}
let mut ret_l = self.action.fold_id();
let mut ret_r = self.action.fold_id();
let mut l = l + self.size / 2;
let mut r = r + self.size / 2;
while l < r {
if r & 1 == 1 {
r -= 1;
self.propagate(r);
ret_r = self.action.fold(self.data[r].clone(), ret_r.clone());
}
if l & 1 == 1 {
self.propagate(l);
ret_l = self.action.fold(ret_l.clone(), self.data[l].clone());
l += 1;
}
r >>= 1;
l >>= 1;
}
self.action.fold(ret_l, ret_r)
}
}
impl<T: Clone + Eq, U: Clone + Eq, A: Clone + Action<FType = T, UType = U>>
Updatable<Range<usize>> for LazySegmentTree<T, U, A>
{
type Value = U;
fn update(&mut self, Range { start: l, end: r }: Range<usize>, x: U) {
self.propagate_top_down(l + self.size / 2);
if r < self.size / 2 {
self.propagate_top_down(r + self.size / 2);
}
{
let mut l = l + self.size / 2;
let mut r = r + self.size / 2;
while l < r {
if r & 1 == 1 {
r -= 1;
self.lazy[r] = self.action.update(x.clone(), self.lazy[r].clone());
self.propagate(r);
}
if l & 1 == 1 {
self.lazy[l] = self.action.update(x.clone(), self.lazy[l].clone());
self.propagate(l);
l += 1;
}
r >>= 1;
l >>= 1;
}
}
self.bottom_up(l + self.size / 2);
if r < self.size / 2 {
self.bottom_up(r + self.size / 2);
}
}
}
}
}
pub mod tree {
pub mod depth_query {
use crate::tree::*;
use std::collections::VecDeque;
pub struct TreeDepthQuery {
par: Vec<Option<usize>>,
depth: Vec<usize>,
left: Vec<usize>,
right: Vec<usize>,
bfs_ord: Vec<Vec<usize>>,
dfs_ord: Vec<Vec<usize>>,
ord: Vec<usize>,
}
impl TreeDepthQuery {
pub fn new<E: TreeEdgeTrait>(tree: &Tree<E>, root: usize) -> Self {
let size = tree.len();
let mut ret = Self {
par: vec![None; size],
depth: vec![0; size],
left: vec![0; size],
right: vec![0; size],
bfs_ord: vec![],
dfs_ord: vec![],
ord: vec![0; size],
};
ret.dfs(&tree, root, None, 0, &mut 0);
let mut q = VecDeque::new();
q.push_back((root, 0));
let mut ord = 0;
while let Some((i, d)) = q.pop_front() {
if ret.bfs_ord.len() <= d {
ret.bfs_ord.push(vec![]);
}
ret.bfs_ord[d].push(ord);
ret.ord[i] = ord;
ord += 1;
for e in tree.nodes[i].neighbors() {
if Some(e.to()) != ret.par[i] {
q.push_back((e.to(), d + 1));
}
}
}
ret
}
fn dfs<E: TreeEdgeTrait>(
&mut self,
tree: &Tree<E>,
cur: usize,
par: Option<usize>,
d: usize,
ord: &mut usize,
) {
self.par[cur] = par;
self.depth[cur] = d;
if self.dfs_ord.len() <= d {
self.dfs_ord.push(vec![]);
}
self.dfs_ord[d].push(*ord);
self.left[cur] = *ord;
*ord += 1;
for e in tree.nodes[cur].neighbors() {
if Some(e.to()) != par {
self.dfs(tree, e.to(), Some(cur), d + 1, ord);
}
}
self.right[cur] = *ord;
}
pub fn children_query(&self, i: usize, d: usize) -> Option<(usize, usize)> {
let d = d + self.depth[i];
if self.bfs_ord.len() > d {
let l = match self.dfs_ord[d].binary_search(&self.left[i]) {
Ok(x) | Err(x) => x,
};
let r = match self.dfs_ord[d].binary_search(&self.right[i]) {
Ok(x) | Err(x) => x,
};
if l >= self.bfs_ord[d].len() {
return None;
}
if r == 0 {
return None;
}
Some((self.bfs_ord[d][l], self.bfs_ord[d][r - 1] + 1))
} else {
None
}
}
pub fn me_query(&self, i: usize) -> (usize, usize) {
(self.ord[i], self.ord[i] + 1)
}
pub fn ancestor(&self, i: usize, k: usize) -> Option<usize> {
let mut p = i;
for _ in 0..k {
match self.par[p] {
Some(x) => p = x,
_ => return None,
}
}
Some(p)
}
}
}
pub trait TreeEdgeTrait {
type Weight;
fn from(&self) -> usize;
fn to(&self) -> usize;
fn weight(&self) -> Self::Weight;
fn rev(self) -> Self;
}
#[derive(Clone, Debug)]
pub struct TreeEdge<T, I> {
pub from: usize,
pub to: usize,
pub weight: T,
pub index: I,
}
impl<T, I> TreeEdge<T, I> {
pub fn new(from: usize, to: usize, weight: T, index: I) -> Self {
Self {
from,
to,
weight,
index,
}
}
}
impl<T: Clone, I> TreeEdgeTrait for TreeEdge<T, I> {
type Weight = T;
#[inline]
fn from(&self) -> usize {
self.from
}
#[inline]
fn to(&self) -> usize {
self.to
}
#[inline]
fn weight(&self) -> Self::Weight {
self.weight.clone()
}
fn rev(mut self) -> Self {
std::mem::swap(&mut self.from, &mut self.to);
self
}
}
#[derive(Clone, Debug)]
pub struct TreeNode<E> {
pub parent: Option<E>,
pub children: Vec<E>,
}
#[derive(Clone, Debug)]
pub struct Tree<E> {
pub nodes: Vec<TreeNode<E>>,
}
impl<E: TreeEdgeTrait> TreeNode<E> {
pub fn neighbors(&self) -> impl DoubleEndedIterator<Item = &E> {
self.children.iter().chain(self.parent.iter())
}
pub fn neighbors_size(&self) -> usize {
self.children.len() + self.parent.as_ref().map_or(0, |_| 1)
}
}
impl<E: TreeEdgeTrait + Clone> Tree<E> {
pub fn new(size: usize) -> Self {
Self {
nodes: vec![
TreeNode {
parent: None,
children: vec![],
};
size
],
}
}
pub fn extend(&mut self, edges: impl IntoIterator<Item = E>) {
for e in edges {
self.nodes[e.from()].children.push(e.clone());
self.nodes[e.to()].children.push(e.rev());
}
}
pub fn extend_rooted(&mut self, edges: impl IntoIterator<Item = E>) {
for e in edges {
assert!(self.nodes[e.to()].parent.is_none());
self.nodes[e.from()].children.push(e.clone());
self.nodes[e.to()].parent.replace(e.rev());
}
}
}
impl<E> Tree<E> {
pub fn len(&self) -> usize {
self.nodes.len()
}
pub fn is_empty(&self) -> bool {
self.nodes.is_empty()
}
}
}
pub mod utils {
pub mod fastio {
use std::fmt::Display;
use std::io::{Read, Write};
pub struct FastIO {
in_bytes: Vec<u8>,
in_cur: usize,
out_buf: std::io::BufWriter<std::io::Stdout>,
}
impl FastIO {
pub fn new() -> Self {
let mut s = vec![];
std::io::stdin().read_to_end(&mut s).unwrap();
let cout = std::io::stdout();
Self {
in_bytes: s,
in_cur: 0,
out_buf: std::io::BufWriter::new(cout),
}
}
#[inline]
pub fn getc(&mut self) -> Option<u8> {
if self.in_cur < self.in_bytes.len() {
self.in_cur += 1;
Some(self.in_bytes[self.in_cur])
} else {
None
}
}
#[inline]
pub fn peek(&self) -> Option<u8> {
if self.in_cur < self.in_bytes.len() {
Some(self.in_bytes[self.in_cur])
} else {
None
}
}
#[inline]
pub fn skip(&mut self) {
while self.peek().map_or(false, |c| c.is_ascii_whitespace()) {
self.in_cur += 1;
}
}
pub fn read_u64(&mut self) -> u64 {
self.skip();
let mut ret: u64 = 0;
while self.peek().map_or(false, |c| c.is_ascii_digit()) {
ret = ret * 10 + (self.in_bytes[self.in_cur] - b'0') as u64;
self.in_cur += 1;
}
ret
}
pub fn read_i64(&mut self) -> i64 {
self.skip();
let mut ret: i64 = 0;
let minus = if self.peek() == Some(b'-') {
self.in_cur += 1;
true
} else {
false
};
while self.peek().map_or(false, |c| c.is_ascii_digit()) {
ret = ret * 10 + (self.in_bytes[self.in_cur] - b'0') as i64;
self.in_cur += 1;
}
if minus {
ret = -ret;
}
ret
}
pub fn read_chars(&mut self) -> Vec<char> {
self.skip();
let mut ret = vec![];
while self.peek().map_or(false, |c| c.is_ascii_graphic()) {
ret.push(self.in_bytes[self.in_cur] as char);
self.in_cur += 1;
}
ret
}
pub fn write<T: Display>(&mut self, s: T) {
self.out_buf.write_all(format!("{}", s).as_bytes()).unwrap();
}
pub fn writeln<T: Display>(&mut self, s: T) {
self.write(s);
self.out_buf.write_all(&[b'\n']).unwrap();
}
}
impl Drop for FastIO {
fn drop(&mut self) {
self.out_buf.flush().unwrap();
}
}
}
}