結果

問題 No.778 クリスマスツリー
ユーザー shino16
提出日時 2021-01-31 19:56:02
言語 Rust
(1.83.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
このコードへのチャレンジ
(要ログイン)
ファイルパターン 結果
sample AC * 3
other AC * 12
権限があれば一括ダウンロードができます

ソースコード

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

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
הההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההה
XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
0