結果
| 問題 |
No.2163 LCA Sum Query
|
| コンテスト | |
| ユーザー |
 akakimidori
|
| 提出日時 | 2022-12-07 04:36:36 |
| 言語 | Rust (1.83.0 + proconio) |
| 結果 |
AC
|
| 実行時間 | 1,218 ms / 6,000 ms |
| コード長 | 18,906 bytes |
| コンパイル時間 | 14,342 ms |
| コンパイル使用メモリ | 388,064 KB |
| 実行使用メモリ | 21,704 KB |
| 最終ジャッジ日時 | 2024-11-06 20:30:56 |
| 合計ジャッジ時間 | 32,163 ms |
|
ジャッジサーバーID （参考情報） |
judge1 / judge4 |
（要ログイン）
| ファイルパターン | 結果 |
|---|---|
| sample | AC * 2 |
| other | AC * 40 |
ソースコード
use std::io::*;
fn main() {
input! {
n: usize,
q: usize,
edge: [(usize, usize); n - 1],
query: [(usize, usize, usize); q],
}
let n = n + 1;
let mut hld = HLD::new(n);
hld.add_edge(0, 1);
for &(a, b) in edge.iter() {
hld.add_edge(a, b);
}
hld.build(0);
let mut cnt = SegmentTreePURQ::new(n, 0, |a, b| *a + *b);
let mut fix_root = LazySegmentTree::new(n, FixRoot);
let mut move_root = LazySegmentTree::new(n, MoveRoot);
for i in 1..n {
let x = hld.vertex(i);
let parent = hld.parent(x).unwrap();
move_root.get_mut(i, |p| p.3 = x as i64 - parent as i64);
}
let mut state = vec![0i64; n];
let mut range = vec![];
let out = std::io::stdout();
let mut out = std::io::BufWriter::new(out.lock());
for &(u, r, v) in query.iter() {
let sign = (state[u] ^ 1) - state[u];
move_root.update(0, n, (sign, 0));
let mut sub = state[u];
hld.path_root(u, &mut range);
for &(l, r) in range.iter().rev() {
move_root.update(l, r, (-sign, sign));
let v = hld.vertex(r - 1);
let p = hld.sub_tree(v);
let c = cnt.find(p.0, p.1) - sub;
fix_root.get_mut(r - 1, |p| {
let v = v as i64;
p.0 += sign * c * v;
p.1 += sign * v;
});
fix_root.update(l, r - 1, sign);
let p = hld.sub_tree(hld.vertex(l));
sub = cnt.find(p.0, p.1);
}
state[u] ^= 1;
cnt.update(hld.sub_tree(u).0, state[u]);
let a = hld.sub_tree(r);
let b = hld.sub_tree(v);
let ans = if b.0 <= a.0 && a.1 <= b.1 {
let mut all = fix_root.find(0, n).0;
hld.path_root(v, &mut range);
for &(l, r) in range.iter() {
all += move_root.find(l, r).0;
}
if r != v {
let x = hld.next(v, r);
let p = hld.sub_tree(x);
let c = cnt.find(0, n);
let d = cnt.find(p.0, p.1);
all -= v as i64 * (c - d) * d;
all -= fix_root.find(p.0, p.1).0;
}
all
} else {
fix_root.find(b.0, b.1).0
};
writeln!(out, "{}", ans).ok();
}
}
struct FixRoot;
impl TE for FixRoot {
type T = (i64, i64);
type E = i64;
fn fold(&self, l: &Self::T, r: &Self::T) -> Self::T {
(l.0 + r.0, l.1 + r.1)
}
fn eval(&self, x: &Self::T, f: &Self::E) -> Self::T {
(x.0 + x.1 * *f, x.1)
}
fn merge(&self, g: &Self::E, h: &Self::E) -> Self::E {
*g + *h
}
fn e(&self) -> Self::T {
(0, 0)
}
fn id(&self) -> Self::E {
0
}
}
struct MoveRoot;
impl TE for MoveRoot {
// sum cab, sum ca, sum cb, sum c
type T = (i64, i64, i64, i64);
// a, b
type E = (i64, i64);
fn fold(&self, l: &Self::T, r: &Self::T) -> Self::T {
(l.0 + r.0, l.1 + r.1, l.2 + r.2, l.3 + r.3)
}
fn eval(&self, x: &Self::T, f: &Self::E) -> Self::T {
(
x.0 + f.0 * x.2 + f.1 * x.1 + f.0 * f.1 * x.3,
x.1 + f.0 * x.3,
x.2 + f.1 * x.3,
x.3,
)
}
fn merge(&self, g: &Self::E, h: &Self::E) -> Self::E {
(g.0 + h.0, g.1 + h.1)
}
fn e(&self) -> Self::T {
(0, 0, 0, 0)
}
fn id(&self) -> Self::E {
(0, 0)
}
}
// ---------- begin segment tree Point Update Range Query ----------
pub struct SegmentTreePURQ<T, F> {
n: usize,
size: usize,
data: Vec<T>,
e: T,
op: F,
}
impl<T, F> SegmentTreePURQ<T, F>
where
T: Clone,
F: Fn(&T, &T) -> T,
{
pub fn new(n: usize, e: T, op: F) -> Self {
assert!(n > 0);
let size = n.next_power_of_two();
let data = vec![e.clone(); 2 * size];
SegmentTreePURQ {
n,
size,
data,
e,
op,
}
}
pub fn update_tmp(&mut self, x: usize, v: T) {
assert!(x < self.n);
self.data[x + self.size] = v;
}
pub fn update_all(&mut self) {
for i in (1..self.size).rev() {
self.data[i] = (self.op)(&self.data[2 * i], &self.data[2 * i + 1]);
}
}
pub fn update(&mut self, x: usize, v: T) {
assert!(x < self.n);
let mut x = x + self.size;
self.data[x] = v;
x >>= 1;
while x > 0 {
self.data[x] = (self.op)(&self.data[2 * x], &self.data[2 * x + 1]);
x >>= 1;
}
}
pub fn find(&self, l: usize, r: usize) -> T {
assert!(l <= r && r <= self.n);
if l == r {
return self.e.clone();
}
let mut l = self.size + l;
let mut r = self.size + r;
let mut x = self.e.clone();
let mut y = self.e.clone();
while l < r {
if l & 1 == 1 {
x = (self.op)(&x, &self.data[l]);
l += 1;
}
if r & 1 == 1 {
r -= 1;
y = (self.op)(&self.data[r], &y);
}
l >>= 1;
r >>= 1;
}
(self.op)(&x, &y)
}
pub fn max_right<P>(&self, l: usize, f: P) -> usize
where
P: Fn(&T) -> bool,
{
assert!(l <= self.n);
assert!(f(&self.e));
if l == self.n {
return self.n;
}
let mut l = l + self.size;
let mut sum = self.e.clone();
while {
l >>= l.trailing_zeros();
let v = (self.op)(&sum, &self.data[l]);
if !f(&v) {
while l < self.size {
l <<= 1;
let v = (self.op)(&sum, &self.data[l]);
if f(&v) {
sum = v;
l += 1;
}
}
return l - self.size;
}
sum = v;
l += 1;
l.count_ones() > 1
} {}
self.n
}
pub fn min_left<P>(&self, r: usize, f: P) -> usize
where
P: Fn(&T) -> bool,
{
assert!(r <= self.n);
assert!(f(&self.e));
if r == 0 {
return 0;
}
let mut r = r + self.size;
let mut sum = self.e.clone();
while {
r -= 1;
while r > 1 && r & 1 == 1 {
r >>= 1;
}
let v = (self.op)(&self.data[r], &sum);
if !f(&v) {
while r < self.size {
r = 2 * r + 1;
let v = (self.op)(&self.data[r], &sum);
if f(&v) {
sum = v;
r -= 1;
}
}
return r + 1 - self.size;
}
sum = v;
(r & (!r + 1)) != r
} {}
0
}
}
// ---------- end segment tree Point Update Range Query ----------
// ---------- begin Heavy-Light decomposition ----------
pub struct HLD {
size: usize,
edge: Vec<(usize, usize)>,
child: Vec<Vec<usize>>,
path_root: Vec<usize>,
parent: Vec<usize>,
left: Vec<usize>,
right: Vec<usize>,
inverse: Vec<usize>,
}
impl HLD {
pub fn new(size: usize) -> Self {
assert!(size <= 10usize.pow(8));
HLD {
size: size,
edge: Vec::with_capacity(size - 1),
child: Vec::new(),
path_root: Vec::new(),
parent: Vec::new(),
left: Vec::new(),
right: Vec::new(),
inverse: Vec::new(),
}
}
pub fn add_edge(&mut self, a: usize, b: usize) {
assert!(a != b && a < self.size && b < self.size);
self.edge.push((a, b));
}
pub fn build(&mut self, root: usize) {
assert!(self.edge.len() + 1 == self.size);
let size = self.size;
let mut cnt = vec![0; size];
for &(a, b) in self.edge.iter() {
cnt[a] += 1;
cnt[b] += 1;
}
let mut child = cnt
.into_iter()
.map(|c| Vec::with_capacity(c))
.collect::<Vec<_>>();
for &(a, b) in self.edge.iter() {
child[a].push(b);
child[b].push(a);
}
let mut parent = vec![size; size];
let mut q = Vec::with_capacity(size);
q.push(root);
parent[root] = root;
for i in 0..size {
let v = q[i];
for u in child[v].clone() {
assert!(parent[u] == size);
parent[u] = v;
child[u].retain(|e| *e != v);
q.push(u);
}
}
let mut sum = vec![1; size];
for &v in q.iter().rev() {
let child = &mut child[v];
if !child.is_empty() {
let (pos, _) = child.iter().enumerate().max_by_key(|p| sum[*p.1]).unwrap();
child.swap(0, pos);
sum[v] = 1 + child.iter().fold(0, |s, a| s + sum[*a]);
}
}
let mut path_root = (0..size).collect::<Vec<_>>();
let mut left = vec![0; size];
let mut right = vec![0; size];
let mut dfs = vec![(root, false)];
let mut id = 0;
while let Some((v, end)) = dfs.pop() {
if end {
right[v] = id;
continue;
}
left[v] = id;
id += 1;
dfs.push((v, true));
let child = &child[v];
if !child.is_empty() {
for &u in child[1..].iter().rev() {
path_root[u] = u;
dfs.push((u, false));
}
let u = child[0];
path_root[u] = path_root[v];
dfs.push((u, false));
}
}
let mut inverse = vec![size; size];
for (i, l) in left.iter().enumerate() {
inverse[*l] = i;
}
self.child = child;
self.parent = parent;
self.left = left;
self.right = right;
self.path_root = path_root;
self.inverse = inverse;
}
pub fn lca(&self, mut a: usize, mut b: usize) -> usize {
assert!(a < self.size && b < self.size);
let path = &self.path_root;
let parent = &self.parent;
let index = &self.left;
while path[a] != path[b] {
if index[a] > index[b] {
std::mem::swap(&mut a, &mut b);
}
b = parent[path[b]];
}
std::cmp::min((index[a], a), (index[b], b)).1
}
pub fn path_root(&self, mut v: usize, res: &mut Vec<(usize, usize)>) {
assert!(v < self.size);
res.clear();
let parent = &self.parent;
let path_root = &self.path_root;
let index = &self.left;
while index[path_root[v]] != 0 {
let p = path_root[v];
res.push((index[p], index[v] + 1));
v = parent[p];
}
res.push((0, index[v] + 1));
res.reverse();
}
pub fn path(
&self,
src: usize,
dst: usize,
up: &mut Vec<(usize, usize)>,
down: &mut Vec<(usize, usize)>,
) {
assert!(src < self.size && dst < self.size);
up.clear();
down.clear();
let path = &self.path_root;
let parent = &self.parent;
let index = &self.left;
let mut x = src;
let mut y = dst;
while path[x] != path[y] {
if index[x] > index[y] {
let p = path[x];
assert!(p == path[p]);
up.push((index[p], index[x] + 1));
x = parent[p];
} else {
let p = path[y];
assert!(p == path[p]);
down.push((index[p], index[y] + 1));
y = parent[p];
}
}
if index[x] >= index[y] {
up.push((index[y], index[x] + 1));
} else {
down.push((index[x], index[y] + 1));
}
down.reverse();
}
pub fn sub_tree(&self, v: usize) -> (usize, usize) {
assert!(v < self.size);
(self.left[v], self.right[v])
}
pub fn parent(&self, v: usize) -> Option<usize> {
assert!(v < self.size);
let p = self.parent[v];
if p == v {
None
} else {
Some(p)
}
}
// s -> t へのパスの2番目の頂点を返す
pub fn next(&self, s: usize, t: usize) -> usize {
assert!(s < self.size && t < self.size && s != t);
let (a, b) = self.sub_tree(s);
let (c, d) = self.sub_tree(t);
if !(a <= c && d <= b) {
return self.parent[s];
}
let mut pos = t;
let mut pre = t;
while self.path_root[s] != self.path_root[pos] {
pre = self.path_root[pos];
pos = self.parent[pre];
}
if s == pos {
pre
} else {
self.child[s][0]
}
}
pub fn vertex(&self, x: usize) -> usize {
assert!(x < self.size);
self.inverse[x]
}
}
// ---------- end Heavy-Light decomposition ----------
// ---------- begin Lazy Segment Tree ----------
pub trait TE {
type T: Clone;
type E: Clone;
fn fold(&self, l: &Self::T, r: &Self::T) -> Self::T;
fn eval(&self, x: &Self::T, f: &Self::E) -> Self::T;
fn merge(&self, g: &Self::E, h: &Self::E) -> Self::E;
fn e(&self) -> Self::T;
fn id(&self) -> Self::E;
}
pub struct LazySegmentTree<R: TE> {
n: usize,
size: usize,
bit: u32,
op: R,
data: Vec<(R::T, R::E)>,
}
impl<R: TE> LazySegmentTree<R> {
pub fn new(n: usize, op: R) -> Self {
assert!(n > 0);
let size = n.next_power_of_two();
let bit = size.trailing_zeros();
let data = vec![(op.e(), op.id()); 2 * size];
Self {
n,
size,
bit,
op,
data,
}
}
pub fn build<I>(init: I, n: usize, op: R) -> Self
where
I: Iterator<Item = R::T>,
{
let mut seg = Self::new(n, op);
for (data, ini) in seg.data[seg.size..].iter_mut().zip(init) {
data.0 = ini;
}
for i in (1..seg.size).rev() {
seg.pull(i);
}
seg
}
pub fn update(&mut self, l: usize, r: usize, f: R::E) {
assert!(l <= r && r <= self.n);
if l == r {
return;
}
self.push_range(l, r);
let mut s = l + self.size;
let mut t = r + self.size;
while s < t {
if s & 1 == 1 {
self.apply(s, &f);
s += 1;
}
if t & 1 == 1 {
t -= 1;
self.apply(t, &f);
}
s >>= 1;
t >>= 1;
}
let l = l + self.size;
let r = r + self.size;
for k in 1..=self.bit {
if (l >> k) << k != l {
self.pull(l >> k);
}
if (r >> k) << k != r {
self.pull((r - 1) >> k);
}
}
}
pub fn find(&mut self, l: usize, r: usize) -> R::T {
assert!(l <= r && r <= self.n);
if l == r {
return self.op.e();
}
self.push_range(l, r);
let mut l = l + self.size;
let mut r = r + self.size;
let mut p = self.op.e();
let mut q = self.op.e();
while l < r {
if l & 1 == 1 {
p = self.op.fold(&p, &self.data[l].0);
l += 1;
}
if r & 1 == 1 {
r -= 1;
q = self.op.fold(&self.data[r].0, &q);
}
l >>= 1;
r >>= 1;
}
self.op.fold(&p, &q)
}
pub fn get_mut<F>(&mut self, x: usize, f: F)
where
F: FnOnce(&mut R::T),
{
assert!(x < self.n);
let x = x + self.size;
for k in (1..=self.bit).rev() {
self.push(x >> k);
}
f(&mut self.data[x].0);
for k in 1..=self.bit {
self.pull(x >> k);
}
}
pub fn set_at(&mut self, x: usize, v: R::T) {
assert!(x < self.n);
let x = x + self.size;
for k in (1..=self.bit).rev() {
self.push(x >> k);
}
self.data[x].0 = v;
for k in 1..=self.bit {
self.pull(x >> k);
}
}
fn push_range(&mut self, l: usize, r: usize) {
let l = l + self.size;
let r = r + self.size;
for k in (1..=self.bit).rev() {
if (l >> k) << k != l {
self.push(l >> k);
}
if (r >> k) << k != r {
self.push((r - 1) >> k);
}
}
}
fn apply(&mut self, x: usize, f: &R::E) {
self.data[x].0 = self.op.eval(&self.data[x].0, f);
self.data[x].1 = self.op.merge(&self.data[x].1, f);
}
fn push(&mut self, x: usize) {
let f = std::mem::replace(&mut self.data[x].1, self.op.id());
self.apply(2 * x, &f);
self.apply(2 * x + 1, &f);
}
fn pull(&mut self, x: usize) {
self.data[x].0 = self.op.fold(&self.data[2 * x].0, &self.data[2 * x + 1].0);
}
}
// ---------- end Lazy Segment Tree ----------
// ---------- begin input macro ----------
// reference: https://qiita.com/tanakh/items/0ba42c7ca36cd29d0ac8
#[macro_export]
macro_rules! input {
(source = $s:expr, $($r:tt)*) => {
let mut iter = $s.split_whitespace();
input_inner!{iter, $($r)*}
};
($($r:tt)*) => {
let s = {
use std::io::Read;
let mut s = String::new();
std::io::stdin().read_to_string(&mut s).unwrap();
s
};
let mut iter = s.split_whitespace();
input_inner!{iter, $($r)*}
};
}
#[macro_export]
macro_rules! input_inner {
($iter:expr) => {};
($iter:expr, ) => {};
($iter:expr, $var:ident : $t:tt $($r:tt)*) => {
let $var = read_value!($iter, $t);
input_inner!{$iter $($r)*}
};
}
#[macro_export]
macro_rules! read_value {
($iter:expr, ( $($t:tt),* )) => {
( $(read_value!($iter, $t)),* )
};
($iter:expr, [ $t:tt ; $len:expr ]) => {
(0..$len).map(|_| read_value!($iter, $t)).collect::<Vec<_>>()
};
($iter:expr, chars) => {
read_value!($iter, String).chars().collect::<Vec<char>>()
};
($iter:expr, bytes) => {
read_value!($iter, String).bytes().collect::<Vec<u8>>()
};
($iter:expr, usize1) => {
read_value!($iter, usize) - 1
};
($iter:expr, $t:ty) => {
$iter.next().unwrap().parse::<$t>().expect("Parse error")
};
}
// ---------- end input macro ----------