#611960 (Rust) No.778 クリスマスツリー

• 問題: No.778 クリスマスツリー
• ユーザー: shino16
• 提出日時: 2021-01-31 19:56:02
• 言語: Rust (1.77.0 + proconio)
• 結果: AC
• 実行時間: 89 ms / 2,000 ms
• コード長: 13,517 bytes
• コンパイル時間: 13,718 ms
• コンパイル使用メモリ: 390,900 KB
• 実行使用メモリ: 21,800 KB
• 最終ジャッジ日時: 2024-06-29 04:18:09
• 合計ジャッジ時間: 12,536 ms
• ジャッジサーバーID （参考情報）: judge2 / judge5

入力	結果	実行時間; 実行使用メモリ
testcase_00	AC	1 ms; 6,812 KB
testcase_01	AC	1 ms; 6,812 KB
testcase_02	AC	1 ms; 6,940 KB
testcase_03	AC	1 ms; 6,944 KB
testcase_04	AC	1 ms; 6,944 KB
testcase_05	AC	1 ms; 6,944 KB
testcase_06	AC	35 ms; 21,140 KB
testcase_07	AC	33 ms; 21,800 KB
testcase_08	AC	78 ms; 20,428 KB
testcase_09	AC	85 ms; 18,400 KB
testcase_10	AC	89 ms; 18,536 KB
testcase_11	AC	56 ms; 18,444 KB
testcase_12	AC	53 ms; 18,312 KB
testcase_13	AC	29 ms; 17,808 KB
testcase_14	AC	30 ms; 21,212 KB

ソースコード

use crate::lib::ds::fenwick::*;
use crate::lib::graph::dfs_io::*;
use crate::lib::io::*;

fn main() {
	let mut io = IO::new();
	let n = io.scan();
	let mut graph = vec![Vec::new(); n];
	for v in 1..n {
		let p: usize = io.scan();
		graph[p].push(v);
		graph[v].push(p);
	}
	let mut fwk = FenwickTree::new(vec![0; n], Addition::new());
	let mut ans = 0;
	dfs_io(&graph, 0, |v, _| match v {
		In(v) => {
			ans += fwk.ask_prefix(v) as u64;
			fwk.add(v, 1_u32);
		},
		Out(v) => fwk.sub(v, 1),
	});
	io.println(ans);
}




pub mod lib {

pub mod alg {


pub trait Alg {
	type Item: Copy;
}

pub trait Monoid: Alg {
	fn unit(&self) -> Self::Item;
	fn op(&self, x: Self::Item, y: Self::Item) -> Self::Item;
	fn op_to(&self, y: Self::Item, x: &mut Self::Item) {
		*x = self.op(*x, y);
	}
}

pub trait Group: Monoid {
	fn inv(&self, x: Self::Item) -> Self::Item;
	fn op_inv_to(&self, y: Self::Item, x: &mut Self::Item) {
		*x = self.op(*x, self.inv(y))
	}
}

macro_rules! impl_monoid {
	($target:ty, $($params:tt : $bounds:tt),*) => {
		impl<$($params : $bounds),*> Alg for $target {
			type Item = T;
		}
		impl<$($params : $bounds),*> Monoid for $target {
			fn unit(&self) -> Self::Item {
				(self.0)()
			}
			fn op(&self, x: Self::Item, y: Self::Item) -> Self::Item {
				(self.1)(x, y)
			}
		}
	};
}

macro_rules! impl_group {
	($target:ty, $($params:tt : $bounds:tt),*) => {
		impl_monoid!($target, $($params : $bounds),*);
		impl<$($params : $bounds),*> Group for $target {
			fn inv(&self, x: Self::Item) -> Self::Item {
				(self.2)(x)
			}
		}
	};
}

pub struct MonoidImpl<T: Copy, Unit: Fn() -> T, Op: Fn(T, T) -> T>(pub Unit, pub Op);
pub struct GroupImpl<T, Unit, Op, Inv>(pub Unit, pub Op, pub Inv)
where
	T: Copy,
	Unit: Fn() -> T,
	Op: Fn(T, T) -> T,
	Inv: Fn(T) -> T;

// help!
impl_monoid!(MonoidImpl<T, Unit, Op>, T: Copy, Unit: (Fn() -> T), Op: (Fn(T, T) -> T));
impl_group!(GroupImpl<T, Unit, Op, Inv>,
			T: Copy, Unit: (Fn() -> T), Op: (Fn(T, T) -> T), Inv: (Fn(T) -> T));

pub mod arith {

pub use super::*;
pub use crate::lib::num::*;
use std::marker::PhantomData;

#[derive(Default)]
pub struct Addition<N>(PhantomData<N>);

impl<N> Addition<N> {
	pub fn new() -> Self {
		Self(PhantomData)
	}
}

impl<N: Num> Alg for Addition<N> {
	type Item = N;
}

impl<N: Num> Monoid for Addition<N> {
	fn unit(&self) -> Self::Item {
		N::ZERO
	}
	fn op(&self, x: Self::Item, y: Self::Item) -> Self::Item {
		x.wrapping_add(y)
	}
}

impl<N: Num> Group for Addition<N> {
	fn inv(&self, x: Self::Item) -> Self::Item {
		x.wrapping_neg()
	}
}

}  // mod arith

}  // mod alg

pub mod bit {

use std::ops::*;

pub trait Bits:
	Sized
	+ BitAnd<Output = Self> + BitAndAssign
	+ BitOr<Output = Self> + BitOrAssign
	+ BitXor<Output = Self> + BitXorAssign
	+ Shl<u32, Output = Self> + ShlAssign<u32>
	+ Shr<u32, Output = Self> + ShrAssign<u32>
	+ Not<Output = Self>
{
	fn trailing_zeros(self) -> u32;
	fn lsb(self) -> Self;
	fn ilog2(self) -> u32;
	fn msb(self) -> Self;
}

macro_rules! impl_bit {
	($($t:ty), *) => { $(
		impl Bits for $t {
			fn trailing_zeros(self) -> u32 {
				<$t>::trailing_zeros(self)
			}
			fn lsb(self) -> Self {
				self & self.wrapping_neg()
			}
			fn ilog2(self) -> u32 {
				std::mem::size_of::<$t>() as u32 * 8 - self.leading_zeros() - 1
			}
			fn msb(self) -> Self {
				(1 as $t) << self.ilog2()
			}
		}
	)* };
}

impl_bit!(i32, i64, i128, isize, u32, u64, u128, usize);

}  // mod bit

pub mod ds {

pub mod bitset {

pub trait BitSet {
	fn get_bit(&self, i: usize) -> bool;
	fn set_bit(&mut self, i: usize, b: bool);
	fn modify_bit(&mut self, i: usize, b: bool) -> bool {
		if self.get_bit(i) == b {
			false
		} else {
			self.set_bit(i, b);
			true
		}
	}
	fn negate(&mut self);
	fn reset(&mut self);
}

macro_rules! impl_bitset {
	($($type:ty),*) => { $(
		impl BitSet for $type {
			fn get_bit(&self, i: usize) -> bool {
				((*self >> i) & 1) != 0
			}
			fn set_bit(&mut self, i: usize, b: bool) {
				*self |= (b as $type) << i;
			}
			fn negate(&mut self) {
				*self = !*self;
			}
			fn reset(&mut self) {
				*self = 0;
			}
		}
	)* };
}

impl_bitset!(i32, i64, i128, isize, u32, u64, u128, usize);

impl BitSet for [u32] {
	fn get_bit(&self, i: usize) -> bool {
		self[i / 32].get_bit(i % 32)
	}
	fn set_bit(&mut self, i: usize, b: bool) {
		self[i / 32].set_bit(i % 32, b);
	}
	fn negate(&mut self) {
		for x in self {
			x.negate()
		}
	}
	fn reset(&mut self) {
		for x in self {
			x.reset();
		}
	}
}

pub fn new_bitset(n: usize) -> Vec<u32> {
	vec![0; (n + 31) / 32]
}

}  // mod bitset

pub mod fenwick {

pub use crate::lib::alg::arith::*;
use crate::lib::bit::*;

#[derive(Clone)]
pub struct FenwickTree<A: Alg> {
	data: Vec<A::Item>,
	alg: A,
}

/// A: Commutative
impl<A: Monoid> FenwickTree<A> {
	pub fn new(mut data: Vec<A::Item>, alg: A) -> Self {
		let len = data.len();
		data.insert(0, alg.unit());
		for i in 1..=len {
			if i + i.lsb() <= len {
				data[i + i.lsb()] = alg.op(data[i + i.lsb()], data[i]);
			}
		}
		Self { data, alg }
	}
	pub fn len(&self) -> usize {
		self.data.len() - 1
	}
	pub fn add(&mut self, pos: usize, v: A::Item) {
		let mut pos = pos + 1;
		while pos < self.data.len() {
			self.data[pos] = self.alg.op(self.data[pos], v);
			pos += pos.lsb();
		}
	}
	pub fn push(&mut self, v: A::Item) {
		self.data.push(self.alg.unit());
		self.add(self.data.len() - 1, v);
	}
	pub fn ask_prefix(&self, mut r: usize) -> A::Item {
		let mut res = self.alg.unit();
		while r != 0 {
			res = self.alg.op(self.data[r], res);
			r -= r.lsb();
		}
		res
	}
	pub fn partition_point<F: FnMut(A::Item) -> bool>(&self, mut pred: F) -> usize {
		let mut x = 0; // pred(&self.ask_prefix(x)) == true
		let mut w = (self.data.len() - 1).msb();
		let mut l = self.alg.unit();
		while w != 0 {
			if x + w < self.data.len() && pred(self.alg.op(l, self.data[x + w])) {
				x += w;
				l = self.alg.op(l, self.data[x + w]);
			}
			w >>= 1;
		}
		x + 1
	}
	pub fn lower_bound(&self, v: A::Item) -> usize
	where
		A::Item: Ord,
	{
		self.partition_point(|x| x < v)
	}
	pub fn upper_bound(&self, v: A::Item) -> usize
	where
		A::Item: Ord,
	{
		self.partition_point(|x| x <= v)
	}
}

/// A: Commutative
impl<A: Group> FenwickTree<A> {
	pub fn sub(&mut self, pos: usize, v: A::Item) {
		self.add(pos, self.alg.inv(v));
	}
	pub fn ask(&self, l: usize, r: usize) -> A::Item {
		self.alg.op(self.alg.inv(self.ask_prefix(l)), self.ask_prefix(r))
	}
}

}  // mod fenwick

}  // mod ds

pub mod graph {



use crate::lib::zo::ZeroOne;

pub trait Graph {
	fn len(&self) -> usize;
	fn adj<F: FnMut(usize)>(&self, v: usize, f: F);
}

pub trait WGraph<W>: Graph {
	fn adj_w<F: FnMut(usize, &W)>(&self, v: usize, f: F);
}

pub trait Tree: Graph {}

pub trait WTree<W>: WGraph<W> {}

impl Graph for Vec<Vec<usize>> {
	fn len(&self) -> usize {
		self.len()
	}
	fn adj<F: FnMut(usize)>(&self, v: usize, f: F) {
		self[v].iter().copied().for_each(f);
	}
}

impl<N: ZeroOne> WGraph<N> for Vec<Vec<usize>> {
	fn adj_w<F: FnMut(usize, &N)>(&self, v: usize, mut f: F) {
		self[v].iter().for_each(|&v| f(v, &N::ONE))
	}
}

impl<W> Graph for Vec<Vec<(usize, W)>> {
	fn len(&self) -> usize {
		self.len()
	}
	fn adj<F: FnMut(usize)>(&self, v: usize, mut f: F) {
		self[v].iter().for_each(|&(v, _)| f(v))
	}
}

impl<W> WGraph<W> for Vec<Vec<(usize, W)>> {
	fn adj_w<F: FnMut(usize, &W)>(&self, v: usize, mut f: F) {
		self[v].iter().for_each(|&(v, ref e)| f(v, e));
	}
}

impl<G: Graph> Tree for G {}
impl<W, G: WGraph<W>> WTree<W> for G {}

pub mod dfs_io {

pub use super::*;
use crate::lib::ds::bitset::*;

#[derive(Debug)]
pub enum InOut {
	In(usize),
	Out(usize),
}

pub use InOut::*;

pub fn dfs_io<G: Graph, F: FnMut(InOut, usize)>(g: &G, s: usize, mut f: F) {
	let mut visited = new_bitset(g.len());
	visited.set_bit(s, true);
	let mut togo = vec![(s, !0)];
	while let Some((v, par)) = togo.pop() {
		if v >> 31 != 0 {
			f(Out(!v), par);
		} else {
			f(In(v), par);
			togo.push((!v, par));
			g.adj(v, |w| {
				if visited.modify_bit(w, true) {
					togo.push((w, v));
				}
			});
		}
	}
}

}  // mod dfs_io

}  // mod graph

pub mod io {

use std::io::{stdout, BufWriter, Read, StdoutLock, Write};
use std::marker::PhantomData;

pub struct IO {
	iter: std::str::SplitAsciiWhitespace<'static>,
	buf: BufWriter<StdoutLock<'static>>,
}

impl IO {
	pub fn new() -> Self {
		let mut input = String::new();
		std::io::stdin().read_to_string(&mut input).unwrap();
		let input = Box::leak(input.into_boxed_str());
		let out = Box::leak(Box::new(stdout()));
		IO {
			iter: input.split_ascii_whitespace(),
			buf: BufWriter::new(out.lock()),
		}
	}
	fn scan_str(&mut self) -> &'static str {
		self.iter.next().unwrap()
	}
	fn scan_raw(&mut self) -> &'static [u8] {
		self.scan_str().as_bytes()
	}
	pub fn scan<T: Scan>(&mut self) -> T {
		T::scan(self)
	}
	pub fn scan_iter<T: Scan>(&mut self, n: usize) -> std::iter::Take<Iter<'_, T>> {
		Iter { io: self, _m: PhantomData }.take(n)
	}
	pub fn scan_vec<T: Scan>(&mut self, n: usize) -> Vec<T> {
		(0..n).map(|_| self.scan()).collect()
	}
	pub fn print<T: Print>(&mut self, x: T) {
		T::print(self, x);
	}
	pub fn println<T: Print>(&mut self, x: T) {
		self.print(x);
		self.print("\n");
	}
	pub fn iterln<T: Print, I: IntoIterator<Item = T>>(&mut self, into_iter: I, delim: &str) {
		let mut iter = into_iter.into_iter();
		if let Some(v) = iter.next() {
			self.print(v);
			for v in iter {
				self.print(delim);
				self.print(v);
			}
		}
		self.print("\n");
	}
	pub fn flush(&mut self) {
		self.buf.flush().unwrap();
	}
}

pub struct Iter<'a, T> {
	io: &'a mut IO,
	_m: PhantomData<T>,
}

impl<T: Scan> Iterator for Iter<'_, T> {
	type Item = T;
	fn next(&mut self) -> Option<Self::Item> {
		Some(self.io.scan())
	}
}

pub trait Scan {
	fn scan(io: &mut IO) -> Self;
}

pub trait Print {
	fn print(w: &mut IO, x: Self);
}

macro_rules! impl_parse_iint {
	($($t:ty),*) => { $(
		impl Scan for $t {
			fn scan(s: &mut IO) -> Self {
				let scan = |t: &[u8]| t.iter().fold(0, |s, &b| s * 10 + (b & 0x0F) as $t);
				let s = s.scan_raw();
				if let Some((&b'-', t)) = s.split_first() { -scan(t) } else { scan(s) }
			}
		}
	)* };
}

macro_rules! impl_parse_uint {
	($($t:ty),*) => { $(
		impl Scan for $t {
			fn scan(s: &mut IO) -> Self {
				s.scan_raw().iter().fold(0, |s, &b| s * 10 + (b & 0x0F) as $t)
			}
		}
	)* };
}

impl_parse_iint!(i32, i64, i128, isize);
impl_parse_uint!(u32, u64, u128, usize);

impl Scan for u8 {
	fn scan(s: &mut IO) -> Self {
		let bytes = s.scan_raw();
		debug_assert_eq!(bytes.len(), 1);
		bytes[0]
	}
}

impl Scan for &[u8] {
	fn scan(s: &mut IO) -> Self {
		s.scan_raw()
	}
}

impl Scan for Vec<u8> {
	fn scan(s: &mut IO) -> Self {
		s.scan_raw().to_vec()
	}
}

macro_rules! impl_tuple {
	() => {};
	($t:ident $($ts:ident)*) => {
		impl<$t: Scan, $($ts: Scan),*> Scan for ($t, $($ts),*) {
			fn scan(s: &mut IO) -> Self { ($t::scan(s), $($ts::scan(s)),*) }
		}
		impl<$t: Print, $($ts: Print),*> Print for ($t, $($ts),*) {
			#[allow(non_snake_case)]
			fn print(w: &mut IO, ($t, $($ts),*): Self) {
				w.print($t);
				$( w.print(" "); w.print($ts); )*
			}
		}
		impl_tuple!($($ts)*);
	};
}

impl_tuple!(A B C D E F G);

macro_rules! impl_scan_array {
	() => {};
	($n:literal $($ns:literal)*) => {
		impl<T: Scan> Scan for [T; $n] {
			fn scan(s: &mut IO) -> Self {
				let mut scan = |_| T::scan(s);
				[scan($n), $(scan($ns)),*]
			}
		}
		impl_scan_array!($($ns)*);
	};
}

impl_scan_array!(7 6 5 4 3 2 1);

macro_rules! impl_print_prim {
	($($t:ty),*) => { $(
		impl Print for $t {
			fn print(w: &mut IO, x: Self) {
				w.buf.write_all(format!("{:.10}", x).as_bytes()).unwrap();
			}
		}
		impl Print for &$t {
			fn print(w: &mut IO, x: Self) { w.print(*x); }
		}
	)* };
}

impl_print_prim!(i32, i64, i128, isize, u32, u64, u128, usize, f32, f64);

impl Print for u8 {
	fn print(w: &mut IO, x: Self) {
		w.buf.write_all(&[x]).unwrap();
	}
}

impl Print for &[u8] {
	fn print(w: &mut IO, x: Self) {
		w.buf.write_all(x).unwrap();
	}
}

impl Print for Vec<u8> {
	fn print(w: &mut IO, x: Self) {
		w.buf.write_all(&x).unwrap();
	}
}

impl Print for &str {
	fn print(w: &mut IO, x: Self) {
		w.print(x.as_bytes());
	}
}

#[derive(Debug, Clone, Copy, Default)]
pub struct Usize1(pub usize);
impl Scan for Usize1 {
	fn scan(io: &mut IO) -> Self {
		let n: usize = io.scan();
		Self(n - 1)
	}
}
impl Print for Usize1 {
	fn print(w: &mut IO, x: Self) {
		w.print(x.0 + 1)
	}
}

}  // mod io

pub mod num {

pub use crate::lib::zo::ZeroOne;
use std::fmt::*;
use std::ops::*;

pub trait Num:
	ZeroOne
	+ Add<Output = Self>
	+ AddAssign
	+ Sub<Output = Self>
	+ SubAssign
	+ Mul<Output = Self>
	+ MulAssign
	+ Div<Output = Self>
	+ DivAssign
	+ Debug
	+ Display
{
	fn wrapping_add(self, rhs: Self) -> Self;
	fn wrapping_neg(self) -> Self;
}

pub trait INum: Num + Neg<Output = Self> {}

macro_rules! impl_num {
	($($t:ty),*) => { $(
		impl Num for $t {
			fn wrapping_add(self, rhs: Self) -> Self {
				self.wrapping_add(rhs)
			}
			fn wrapping_neg(self) -> Self {
				self.wrapping_neg()
			}
		}
	)* };
}

impl_num!(i32, i64, i128, isize, u32, u64, u128, usize);

impl<T: Num + Neg<Output = Self>> INum for T {}

}  // mod num

pub mod zo {

pub trait ZeroOne: Copy + Eq {
	const ZERO: Self;
	fn is_zero(self) -> bool {
		self == Self::ZERO
	}
	const ONE: Self;
}

macro_rules! impl_zo {
	($($t:ty),*) => { $(
		impl ZeroOne for $t {
			const ZERO: Self = 0;
			const ONE: Self = 1;
		}
	)* };
}

impl_zo!(i32, i64, i128, isize, u32, u64, u128, usize);

}  // mod zo

}  // mod lib