結果
| 問題 |
No.399 動的な領主
|
| コンテスト | |
| ユーザー |
ngtkana
|
| 提出日時 | 2024-06-24 02:26:20 |
| 言語 | Rust (1.83.0 + proconio) |
| 結果 |
AC
|
| 実行時間 | 884 ms / 2,000 ms |
| コード長 | 24,146 bytes |
| コンパイル時間 | 16,984 ms |
| コンパイル使用メモリ | 400,916 KB |
| 実行使用メモリ | 30,564 KB |
| 最終ジャッジ日時 | 2024-06-24 02:26:49 |
| 合計ジャッジ時間 | 23,083 ms |
|
ジャッジサーバーID (参考情報) |
judge1 / judge5 |
(要ログイン)
| ファイルパターン | 結果 |
|---|---|
| other | AC * 19 |
ソースコード
use lazy_segtree::LazySegtree;
use proconio::input;
use proconio::marker::Usize1;
fn main() {
input! {
n: usize,
edges: [(Usize1, Usize1); n - 1],
q: usize,
queries: [(Usize1, Usize1); q],
}
let (hld, _g) = Hld::from_edges(0, &edges);
let mut lazy_segtree = LazySegtree::<O>::from_iter(vec![Value { len: 1, sum: 1 }; n]);
let mut ans = 0;
for &(i, j) in &queries {
hld.visit_path_segments_including_lca_by_index(i, j, |i, j| {
ans += lazy_segtree.fold(i..=j).sum;
lazy_segtree.range_apply(i..=j, &1);
});
}
println!("{ans}");
}
#[derive(Clone, Copy, Debug, PartialEq, Default)]
struct Value {
len: usize,
sum: usize,
}
enum O {}
impl lazy_segtree::Op for O {
type Operator = usize;
type Value = Value;
fn identity() -> Self::Value {
Value { len: 0, sum: 0 }
}
fn op(lhs: &Self::Value, rhs: &Self::Value) -> Self::Value {
Value {
len: lhs.len + rhs.len,
sum: lhs.sum + rhs.sum,
}
}
fn apply(op: &Self::Operator, value: &Self::Value) -> Self::Value {
Value {
len: value.len,
sum: value.sum + op * value.len,
}
}
fn identity_op() -> Self::Operator {
0
}
fn compose(op: &Self::Operator, other: &Self::Operator) -> Self::Operator {
op + other
}
}
pub struct Hld {
pub parent: Vec<usize>,
pub index: Vec<usize>,
pub head: Vec<usize>,
}
impl Hld {
pub fn from_short_parents(mut parent: Vec<usize>) -> (Self, Vec<Vec<usize>>) {
parent.insert(0, 0);
let mut g = vec![Vec::new(); parent.len()];
for (i, &p) in parent.iter().enumerate().skip(1) {
g[p].push(i);
}
(__build_hld(0, &mut g, parent), g)
}
pub fn from_edges(root: usize, edges: &[(usize, usize)]) -> (Self, Vec<Vec<usize>>) {
let mut g = vec![Vec::new(); edges.len() + 1];
for &(i, j) in edges {
g[i].push(j);
g[j].push(i);
}
let parent = __remove_parent(root, &mut g);
(__build_hld(root, &mut g, parent), g)
}
pub fn visit_path_segments_including_lca(
&self,
mut i: usize, // id
mut j: usize, // id
mut f: impl FnMut(usize, usize), // id
) {
while self.head[i] != self.head[j] {
if self.index[i] < self.index[j] {
f(self.head[j], j);
j = self.parent[self.head[j]];
} else {
f(self.head[i], i);
i = self.parent[self.head[i]];
}
}
if self.index[i] < self.index[j] {
f(i, j)
} else {
f(j, i)
}
}
pub fn visit_path_segments_including_lca_by_index(
&self,
i: usize, // id
j: usize, // id
mut f: impl FnMut(usize, usize), // index
) {
self.visit_path_segments_including_lca(i, j, |i, j| f(self.index[i], self.index[j]));
}
pub fn lca(&self, mut i: usize, mut j: usize) -> usize {
while self.head[i] != self.head[j] {
if self.index[i] < self.index[j] {
j = self.parent[self.head[j]];
} else {
i = self.parent[self.head[i]];
}
}
std::cmp::min_by_key(i, j, |&i| self.index[i])
}
}
fn __build_hld(root: usize, g: &mut [Vec<usize>], parent: Vec<usize>) -> Hld {
let n = g.len();
__heavy_first(0, g);
let mut index = vec![usize::MAX; n];
let mut head = vec![usize::MAX; n];
head[root] = root;
__head_and_index(0, &*g, &mut head, &mut index, &mut (0..));
Hld {
parent,
index,
head,
}
}
fn __head_and_index(
i: usize,
g: &[Vec<usize>],
head: &mut [usize],
index: &mut [usize],
current: &mut std::ops::RangeFrom<usize>,
) {
index[i] = current.next().unwrap();
for &j in &g[i] {
head[j] = if j == g[i][0] { head[i] } else { j };
__head_and_index(j, g, head, index, current);
}
}
fn __heavy_first(i: usize, g: &mut [Vec<usize>]) -> usize {
let mut max = 0;
let mut size = 1;
for e in 0..g[i].len() {
let csize = __heavy_first(g[i][e], g);
if max < csize {
max = csize;
g[i].swap(0, e);
}
size += csize;
}
size
}
fn __remove_parent(root: usize, g: &mut [Vec<usize>]) -> Vec<usize> {
let mut stack = vec![root];
let mut parent = vec![usize::MAX; g.len()];
parent[root] = root;
while let Some(i) = stack.pop() {
g[i].retain(|&j| parent[i] != j);
for &j in &g[i] {
parent[j] = i;
stack.push(j);
}
}
parent
}
// link_cut_tree {{{
// https://ngtkana.github.io/ac-adapter-rs/link_cut_tree/index.html
#[allow(dead_code)]
mod link_cut_tree {
mod base {
#[doc(hidden)]
pub trait OpBase {
type Value: Clone;
type InternalValue: Clone;
fn identity() -> Self::InternalValue;
fn mul(lhs: &Self::InternalValue, rhs: &Self::InternalValue) -> Self::InternalValue;
fn into_front(value: Self::InternalValue) -> Self::Value;
fn from_front(value: Self::Value) -> Self::InternalValue;
fn rev(value: &mut Self::InternalValue);
}
pub struct LinkCutTreeBase<O: OpBase> {
nodes: Vec<Node<O>>,
}
impl<O: OpBase> LinkCutTreeBase<O> {
pub fn new(n: usize) -> Self {
Self {
nodes: (0..n)
.map(|id| Node {
id,
parent: std::ptr::null_mut(),
left: std::ptr::null_mut(),
right: std::ptr::null_mut(),
rev: false,
value: O::identity(),
acc: O::identity(),
})
.collect(),
}
}
pub fn from_values(values: impl IntoIterator<Item = O::Value>) -> Self {
Self {
nodes: values
.into_iter()
.map(O::from_front)
.enumerate()
.map(|(id, value)| Node {
id,
parent: std::ptr::null_mut(),
left: std::ptr::null_mut(),
right: std::ptr::null_mut(),
rev: false,
value: value.clone(),
acc: value,
})
.collect(),
}
}
pub fn link(&mut self, p: usize, c: usize) {
unsafe {
let c = std::ptr::addr_of_mut!(self.nodes[c]);
let p = std::ptr::addr_of_mut!(self.nodes[p]);
expose(c);
assert!((*c).left.is_null(), "c = {} is not a root", (*c).id);
expose(p);
assert!(
(*c).parent.is_null(),
"c = {} and p = {} are already connected",
(*c).id,
(*p).id
);
(*c).parent = p;
(*p).right = c;
update(p);
}
}
pub fn undirected_link(&mut self, i: usize, j: usize) -> bool {
if self.undirected_is_connected(i, j) {
return false;
}
self.evert(j);
self.link(i, j);
true
}
pub fn cut(&mut self, x: usize) -> Option<usize> {
unsafe {
let x = std::ptr::addr_of_mut!(self.nodes[x]);
expose(x);
let p = (*x).left;
(*x).left = std::ptr::null_mut();
let ans = p.as_ref().map(|p| p.id);
if !p.is_null() {
(*p).parent = std::ptr::null_mut();
}
update(x);
ans
}
}
pub fn undirected_cut(&mut self, i: usize, j: usize) -> bool {
if !self.undirected_has_edge(i, j) {
return false;
}
self.evert(i);
self.cut(j);
true
}
pub fn evert(&mut self, x: usize) {
unsafe {
let x = std::ptr::addr_of_mut!(self.nodes[x]);
expose(x);
rev(x);
push(x);
}
}
pub fn undirected_has_edge(&mut self, x: usize, y: usize) -> bool {
self.parent(x) == Some(y) || self.parent(y) == Some(x)
}
pub fn undirected_is_connected(&mut self, x: usize, y: usize) -> bool {
if x == y {
return true;
}
unsafe {
let x = std::ptr::addr_of_mut!(self.nodes[x]);
let y = std::ptr::addr_of_mut!(self.nodes[y]);
expose(x);
expose(y);
!(*x).parent.is_null()
}
}
pub fn lca(&mut self, x: usize, y: usize) -> Option<usize> {
if x == y {
return Some(x);
}
unsafe {
let x = std::ptr::addr_of_mut!(self.nodes[x]);
let y = std::ptr::addr_of_mut!(self.nodes[y]);
expose(x);
let lca = expose(y);
if (*x).parent.is_null() {
None
} else {
Some((*lca).id)
}
}
}
pub fn set(&mut self, x: usize, mut f: impl FnMut(O::Value) -> O::Value) {
unsafe {
let x = std::ptr::addr_of_mut!(self.nodes[x]);
expose(x);
(*x).value = O::from_front(f(O::into_front((*x).value.clone())));
update(x);
}
}
pub fn fold(&mut self, x: usize) -> O::Value {
unsafe {
let x = std::ptr::addr_of_mut!(self.nodes[x]);
expose(x);
O::into_front((*x).acc.clone())
}
}
pub fn undirected_fold(&mut self, i: usize, j: usize) -> Option<O::Value> {
if !self.undirected_is_connected(i, j) {
return None;
}
self.evert(i);
Some(self.fold(j))
}
pub fn parent(&mut self, x: usize) -> Option<usize> {
unsafe {
let x = std::ptr::addr_of_mut!(self.nodes[x]);
expose(x);
let mut p = (*x).left.as_mut()?;
while let Some(next) = p.right.as_mut() {
p = next;
}
splay(p);
Some(p.id)
}
}
}
#[derive(Clone, Copy)]
struct Node<O: OpBase> {
id: usize,
parent: *mut Self,
left: *mut Self,
right: *mut Self,
rev: bool,
value: O::InternalValue,
acc: O::InternalValue,
}
unsafe fn is_splay_root<O: OpBase>(x: *mut Node<O>) -> bool {
let x = &*x;
let p = match x.parent.as_ref() {
Some(p) => p,
None => return true,
};
!std::ptr::eq(x, p.left) && !std::ptr::eq(x, p.right)
}
unsafe fn push<O: OpBase>(x: *mut Node<O>) {
let x = &mut *x;
if x.rev {
if let Some(l) = x.left.as_mut() {
rev(l);
}
if let Some(r) = x.right.as_mut() {
rev(r);
}
x.rev = false;
}
}
unsafe fn update<O: OpBase>(x: *mut Node<O>) {
let x = &mut *x;
x.acc = x.value.clone();
if !x.left.is_null() {
x.acc = O::mul(&(*x.left).acc, &x.acc);
}
if !x.right.is_null() {
x.acc = O::mul(&x.acc, &(*x.right).acc);
}
}
unsafe fn rev<O: OpBase>(x: *mut Node<O>) {
let x = &mut *x;
std::mem::swap(&mut x.left, &mut x.right);
O::rev(&mut x.acc);
x.rev ^= true;
}
unsafe fn expose<O: OpBase>(x: *mut Node<O>) -> *mut Node<O> {
let mut last = std::ptr::null_mut();
let mut current = x;
while !current.is_null() {
splay(current);
(*current).right = last;
update(current);
last = current;
current = (*current).parent;
}
splay(x);
last
}
unsafe fn splay<O: OpBase>(x: *mut Node<O>) {
let x = &mut *x;
push(x);
while !is_splay_root(x) {
let p = &mut *x.parent;
if is_splay_root(p) {
push(p);
push(x);
if std::ptr::eq(p.left, x) {
rotate_right(p);
} else {
rotate_left(p);
}
} else {
let g = &mut *p.parent;
push(g);
push(p);
push(x);
#[allow(clippy::collapsible_else_if)]
if std::ptr::eq(p.left, x) {
if std::ptr::eq(g.left, p) {
rotate_right(g);
rotate_right(p);
} else {
rotate_right(p);
rotate_left(g);
}
} else {
if std::ptr::eq(g.left, p) {
rotate_left(p);
rotate_right(g);
} else {
rotate_left(g);
rotate_left(p);
}
}
}
}
}
unsafe fn rotate_left<O: OpBase>(l: *mut Node<O>) {
let l = &mut *l;
let r = &mut *l.right;
let p = l.parent;
let c = r.left;
l.right = c;
if !c.is_null() {
(*c).parent = l;
}
r.left = l;
l.parent = r;
r.parent = p;
update(l);
update(r);
if !p.is_null() {
if std::ptr::eq((*p).left, l) {
(*p).left = r;
} else if std::ptr::eq((*p).right, l) {
(*p).right = r;
}
update(&mut *p);
}
}
unsafe fn rotate_right<O: OpBase>(r: *mut Node<O>) {
let r = &mut *r;
let l = &mut *r.left;
let p = r.parent;
let c = l.right;
r.left = c;
if !c.is_null() {
(*c).parent = r;
}
l.right = r;
r.parent = l;
l.parent = p;
update(r);
update(l);
if !p.is_null() {
if std::ptr::eq((*p).left, r) {
(*p).left = l;
} else if std::ptr::eq((*p).right, r) {
(*p).right = l;
}
update(&mut *p);
}
}
}
pub use base::LinkCutTreeBase;
use base::OpBase;
pub trait Op {
type Value: Clone;
fn identity() -> Self::Value;
fn mul(lhs: &Self::Value, rhs: &Self::Value) -> Self::Value;
}
impl OpBase for () {
type InternalValue = ();
type Value = ();
fn identity() -> Self::InternalValue {}
fn mul(_lhs: &Self::InternalValue, _rhs: &Self::InternalValue) -> Self::InternalValue {}
fn rev(_value: &mut Self::InternalValue) {}
fn into_front(_value: Self::InternalValue) {}
fn from_front(_value: Self::Value) -> Self::InternalValue {}
}
pub type LinkCutTree = LinkCutTreeBase<()>;
pub type CommutLinkCutTree<T> = LinkCutTreeBase<Commut<T>>;
#[doc(hidden)]
pub struct Commut<T: Op>(T);
impl<T: Op> OpBase for Commut<T> {
type InternalValue = T::Value;
type Value = T::Value;
fn identity() -> Self::InternalValue {
T::identity()
}
fn mul(lhs: &Self::InternalValue, rhs: &Self::InternalValue) -> Self::InternalValue {
T::mul(lhs, rhs)
}
fn rev(_value: &mut Self::InternalValue) {}
fn into_front(value: Self::InternalValue) -> Self::Value {
value
}
fn from_front(value: Self::Value) -> Self::InternalValue {
value
}
}
#[doc(hidden)]
pub struct NonCommut<T: Op>(T);
pub type NonCommutLinkCutTree<T> = LinkCutTreeBase<NonCommut<T>>;
impl<T: Op> OpBase for NonCommut<T> {
type InternalValue = (T::Value, T::Value);
type Value = T::Value;
fn identity() -> Self::InternalValue {
(T::identity(), T::identity())
}
fn mul(lhs: &Self::InternalValue, rhs: &Self::InternalValue) -> Self::InternalValue {
(T::mul(&lhs.0, &rhs.0), T::mul(&rhs.1, &lhs.1))
}
fn rev(value: &mut Self::InternalValue) {
std::mem::swap(&mut value.0, &mut value.1);
}
fn into_front(value: Self::InternalValue) -> Self::Value {
value.0
}
fn from_front(value: Self::Value) -> Self::InternalValue {
(value.clone(), value)
}
}
}
// }}}
// lazy_segtree {{{
// https://ngtkana.github.io/ac-adapter-rs/lazy_segtree/index.html
#[allow(dead_code)]
mod lazy_segtree {
use std::iter::FromIterator;
use std::mem::replace;
use std::ops::RangeBounds;
pub trait Op {
type Value;
type Operator: PartialEq;
fn identity() -> Self::Value;
fn op(lhs: &Self::Value, rhs: &Self::Value) -> Self::Value;
fn apply(op: &Self::Operator, value: &Self::Value) -> Self::Value;
fn identity_op() -> Self::Operator;
fn compose(op: &Self::Operator, other: &Self::Operator) -> Self::Operator;
}
#[derive(Debug, Clone)]
pub struct LazySegtree<O: Op> {
values: Vec<O::Value>,
operators: Vec<O::Operator>,
}
impl<O: Op> LazySegtree<O> {
pub fn new(values: &[O::Value]) -> Self
where
O::Value: Clone,
O::Operator: Clone,
{
let values_ = values;
let n = values_.len();
let mut values = vec![O::identity(); 2 * n];
values[n..].clone_from_slice(values_);
for i in (1..n).rev() {
values[i] = O::op(&values[i * 2], &values[i * 2 + 1]);
}
Self {
values,
operators: vec![O::identity_op(); 2 * n],
}
}
pub fn range_apply<R: RangeBounds<usize>>(&mut self, range: R, f: &O::Operator) {
let n = self.operators.len() / 2;
let (l, r) = open(range, n);
if l == r {
return;
}
let l = n + l;
let r = n + r;
for p in (0..usize::BITS - r.leading_zeros()).rev() {
self.push(l >> p);
self.push((r - 1) >> p);
}
{
let mut l = l;
let mut r = r;
while l < r {
if l & 1 != 0 {
self.operators[l] = O::compose(f, &self.operators[l]);
l += 1;
}
if r & 1 != 0 {
r -= 1;
self.operators[r] = O::compose(f, &self.operators[r]);
}
l >>= 1;
r >>= 1;
}
}
for p in 1..usize::BITS - r.leading_zeros() {
self.update(l >> p);
self.update((r - 1) >> p);
}
}
pub fn fold<R: RangeBounds<usize>>(&mut self, range: R) -> O::Value {
let n = self.operators.len() / 2;
let (mut l, mut r) = open(range, n);
if l == r {
return O::identity();
}
l += n;
r += n;
for p in (0..usize::BITS - r.leading_zeros()).rev() {
self.push(l >> p);
self.push((r - 1) >> p);
}
for p in 1..usize::BITS - r.leading_zeros() {
self.update(l >> p);
self.update((r - 1) >> p);
}
let mut left = O::identity();
let mut right = O::identity();
while l < r {
if l & 1 != 0 {
left = O::op(&left, &O::apply(&self.operators[l], &self.values[l]));
l += 1;
}
if r & 1 != 0 {
r -= 1;
right = O::op(&O::apply(&self.operators[r], &self.values[r]), &right);
}
l >>= 1;
r >>= 1;
}
O::op(&left, &right)
}
pub fn get(&self, i: usize) -> O::Value
where
O::Value: Clone,
{
let mut i = self.operators.len() / 2 + i;
let mut value = self.values[i].clone();
while i > 0 {
value = O::apply(&self.operators[i], &value);
i >>= 1;
}
value
}
pub fn collect(&self) -> Vec<O::Value>
where
O::Value: Clone,
{
(0..self.operators.len() / 2).map(|i| self.get(i)).collect()
}
fn push(&mut self, i: usize) {
let a = replace(&mut self.operators[i], O::identity_op());
self.values[i] = O::apply(&a, &self.values[i]);
if i < self.operators.len() / 2 {
self.operators[i << 1] = O::compose(&a, &self.operators[i << 1]);
self.operators[i << 1 | 1] = O::compose(&a, &self.operators[i << 1 | 1]);
}
}
fn update(&mut self, i: usize) {
self.values[i] = O::op(
&O::apply(&self.operators[i << 1], &self.values[i << 1]),
&O::apply(&self.operators[i << 1 | 1], &self.values[i << 1 | 1]),
)
}
}
impl<O: Op> FromIterator<O::Value> for LazySegtree<O>
where
O::Value: Clone,
O::Operator: Clone,
{
fn from_iter<T: IntoIterator<Item = O::Value>>(iter: T) -> Self {
Self::new(&iter.into_iter().collect::<Vec<_>>())
}
}
fn open<B: RangeBounds<usize>>(bounds: B, n: usize) -> (usize, usize) {
use std::ops::Bound;
let start = match bounds.start_bound() {
Bound::Unbounded => 0,
Bound::Included(&x) => x,
Bound::Excluded(&x) => x + 1,
};
let end = match bounds.end_bound() {
Bound::Unbounded => n,
Bound::Included(&x) => x + 1,
Bound::Excluded(&x) => x,
};
(start, end)
}
}
// }}}
ngtkana