結果
| 問題 |
No.899 γatheree
|
| コンテスト | |
| ユーザー |
 Haar
|
| 提出日時 | 2021-12-01 23:30:37 |
| 言語 | Rust (1.83.0 + proconio) |
| 結果 |
RE
|
| 実行時間 | - |
| コード長 | 17,995 bytes |
| コンパイル時間 | 13,558 ms |
| コンパイル使用メモリ | 402,656 KB |
| 実行使用メモリ | 25,932 KB |
| 最終ジャッジ日時 | 2024-07-05 01:30:19 |
| 合計ジャッジ時間 | 22,164 ms |
|
ジャッジサーバーID （参考情報） |
judge3 / judge5 |
（要ログイン）
| ファイルパターン | 結果 |
|---|---|
| sample | AC * 1 |
| other | AC * 3 RE * 20 |
ソースコード
#![allow(unused_imports)]
use std::io::{Read, Write};
#[allow(unused_macros)]
macro_rules! get {
( $in:ident, [$a:tt; $num:expr] ) => {
{
let n = $num;
(0 .. n).map(|_| get!($in, $a)).collect::<Vec<_>>()
}
};
( $in:ident, ($($type:ty),*) ) => {
($(get!($in, $type)),*)
};
( $in:ident, $type:ty ) => {
{
let token = $in.next().unwrap();
token.parse::<$type>().expect(
format!("cannot convert \"{}\" into {}", token, stringify!($type)).as_str()
)
}
};
}
macro_rules! input {
( @inner $in:ident, mut $name:ident : $type:tt ) => {
let mut $name = get!($in, $type);
};
( @inner $in:ident, $name:ident : $type:tt ) => {
let $name = get!($in, $type);
};
( $in:ident, $($($names:ident)* : $type:tt),* ) => {
$(
input!(@inner $in, $($names)* : $type);
)*
}
}
#[allow(unused_macros)]
macro_rules! io {
( $in:ident, $out:ident ) => {
let mut s = String::new();
std::io::stdin().read_to_string(&mut s).unwrap();
let mut $in = s.split_ascii_whitespace();
let $out = std::io::stdout();
let mut $out = std::io::BufWriter::new($out.lock());
};
}
pub mod main {
use super::*;
use haar_lib::{
algebra::update_sum::*,
ds::lazy_segtree::*,
//chmin, chmax,
//mul_vec,
tree::{depth_query::*, *},
};
#[derive(Clone, Default)]
pub struct Problem {/* write variables here */}
impl Problem {
pub fn main(&mut self) -> Result<(), Box<dyn std::error::Error>> {
io!(cin, cout);
input!(cin, n: usize);
let mut tree = Tree::new(n);
for _ in 0..n - 1 {
input!(cin, u: usize, v: usize);
tree.add_undirected(vec![(u, v, ())]);
}
let res = TreeDepthQuery::new(&tree, 0);
input!(cin, a: [u64; n]);
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]));
}
input!(cin, q: usize);
for _ in 0..q {
input!(cin, x: 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).map(|p| f(p));
}
f(res.me_query(x));
res.children_query(x, 1).map(|p| f(p));
res.children_query(x, 2).map(|p| f(p));
let (l, r) = res.me_query(x);
seg.update(l..r, Some(ans));
writeln!(cout, "{}", ans)?;
}
Ok(())
}
/* write functions here */
}
}
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 algo {
pub mod bsearch {
#[doc = " x以上となる最小のindexを求める。"]
#[doc = ""]
#[doc = " # Complexity"]
#[doc = " Time complexity $O(\\log(n))$"]
pub fn lower_bound<T: Ord>(a: &[T], value: &T) -> usize {
let n = a.len();
let mut lb = 0;
let mut len = n;
while len > 0 {
let half = len / 2;
let mid = lb + half;
if &a[mid] < value {
len -= half + 1;
lb = mid + 1;
} else {
len = half;
}
}
lb
}
#[doc = " xを超える最小のindexを求める。"]
#[doc = ""]
#[doc = " # Complexity"]
#[doc = " Time complexity $O(\\log(n))$"]
pub fn upper_bound<T: Ord>(a: &[T], value: &T) -> usize {
let n = a.len();
let mut ub = 0;
let mut len = n;
while len > 0 {
let half = len / 2;
let mid = ub + half;
if &a[mid] <= value {
len -= half + 1;
ub = mid + 1;
} else {
len = half;
}
}
ub
}
#[doc = " lower_bound, upper_boundの組を求める。"]
#[doc = ""]
#[doc = " # Complexity"]
#[doc = " Time complexity $O(\\log(n))$"]
pub fn equal_range<T: Ord>(a: &[T], value: &T) -> (usize, usize) {
(lower_bound(a, value), upper_bound(a, value))
}
}
}
pub mod ds {
pub mod traits {
pub trait Foldable<Idx> {
type Output;
fn fold(&self, range: Idx) -> Self::Output;
}
pub trait FoldableMut<Idx> {
type Output;
fn fold(&mut self, 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 IndexableMut<Idx> {
type Output;
fn get(&mut self, i: Idx) -> Self::Output;
}
}
pub mod lazy_segtree {
use crate::algebra::action::Action;
pub use crate::ds::traits::*;
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());
}
}
}
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);
}
}
}
impl<T: Clone + Eq, U: Clone + Eq, A: Clone + Action<FType = T, UType = U>>
FoldableMut<Range<usize>> for LazySegmentTree<T, U, A>
{
type Output = T;
fn fold(&mut self, Range { start: l, end: r }: Range<usize>) -> Self::Output {
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)
}
}
}
}
pub mod tree {
pub mod depth_query {
use crate::{algo::bsearch::lower_bound, 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<T>(tree: &Tree<T>, 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 &TreeEdge { to, .. } in tree.nodes[i].neighbors() {
if Some(to) != ret.par[i] {
q.push_back((to, d + 1));
}
}
}
ret
}
fn dfs<T>(
&mut self,
tree: &Tree<T>,
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 &TreeEdge { to, .. } in tree.nodes[cur].neighbors() {
if Some(to) != par {
self.dfs(tree, 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 = lower_bound(&self.dfs_ord[d], &self.left[i]);
let r = lower_bound(&self.dfs_ord[d], &self.right[i]);
if l >= self.bfs_ord[d].len() {
return None;
}
if r == 1 {
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)
}
}
}
#[derive(Clone, Debug)]
pub struct TreeEdge<T> {
pub to: usize,
pub weight: T,
}
#[derive(Clone, Debug)]
pub struct TreeNode<T> {
pub index: usize,
pub parent: Option<TreeEdge<T>>,
pub children: Vec<TreeEdge<T>>,
}
#[derive(Clone, Debug)]
pub struct Tree<T> {
pub nodes: Vec<TreeNode<T>>,
}
impl<T> TreeNode<T> {
pub fn neighbors(&self) -> impl DoubleEndedIterator<Item = &TreeEdge<T>> {
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<T: Copy> Tree<T> {
pub fn new(size: usize) -> Self {
Self {
nodes: (0..size)
.map(|i| TreeNode {
index: i,
parent: None,
children: vec![],
})
.collect(),
}
}
pub fn add_undirected(&mut self, edges: impl IntoIterator<Item = (usize, usize, T)>) {
for (u, v, w) in edges {
self.nodes[u].children.push(TreeEdge { to: v, weight: w });
self.nodes[v].children.push(TreeEdge { to: u, weight: w });
}
}
pub fn add_directed(&mut self, edges: impl IntoIterator<Item = (usize, usize, T)>) {
for (p, c, w) in edges {
assert!(self.nodes[c].parent.is_none());
self.nodes[p].children.push(TreeEdge { to: c, weight: w });
self.nodes[c].parent = Some(TreeEdge { to: p, weight: w });
}
}
}
impl<T> Tree<T> {
pub fn len(&self) -> usize {
self.nodes.len()
}
pub fn is_empty(&self) -> bool {
self.nodes.is_empty()
}
}
}