結果
| 問題 |
No.3194 Do Optimize Your Solution
|
| ユーザー |
 akakimidori
|
| 提出日時 | 2025-06-29 04:43:27 |
| 言語 | Rust (1.83.0 + proconio) |
| 結果 |
TLE
|
| 実行時間 | - |
| コード長 | 14,297 bytes |
| コンパイル時間 | 14,495 ms |
| コンパイル使用メモリ | 398,792 KB |
| 実行使用メモリ | 85,216 KB |
| 最終ジャッジ日時 | 2025-06-29 04:44:08 |
| 合計ジャッジ時間 | 22,376 ms |
|
ジャッジサーバーID （参考情報） |
judge1 / judge5 |
（要ログイン）
| ファイルパターン | 結果 |
|---|---|
| sample | AC * 2 |
| other | AC * 16 TLE * 1 |
コンパイルメッセージ
warning: function `sack` is never used
--> src/main.rs:49:4
|
49 | fn sack(v: usize, save: bool, g: &[Vec<usize>], stt: &mut STT<State>, ans: &mut u64) {
| ^^^^
|
= note: `#[warn(dead_code)]` on by default
warning: function `flip` is never used
--> src/main.rs:64:4
|
64 | fn flip(v: usize, g: &[Vec<usize>], stt: &mut STT<State>) {
| ^^^^
ソースコード
fn main() {
input! {
n: usize,
a: [(usize1, usize1); n - 1],
b: [(usize1, usize1); n - 1],
}
let mut g = vec![vec![]; n];
for (a, b) in a {
g[a].push(b);
g[b].push(a);
}
hld(&mut g);
let state = State { vertex: vec![0; n] };
let mut stt = STT::new(state, 0, b);
let mut ans = 0;
let mut dfs = vec![(0, true, false)];
let mut memo = vec![];
while let Some((v, save, eval)) = dfs.pop() {
if eval {
for &u in g[v].iter().skip(1) {
memo.push(u);
}
while let Some(u) = memo.pop() {
stt.state.vertex[u] ^= 1;
stt.update_vertex(u);
memo.extend(g[u].iter().cloned());
}
stt.state.vertex[v] ^= 1;
stt.update_vertex(v);
ans += stt.find().ans;
if !save {
memo.push(v);
while let Some(u) = memo.pop() {
stt.state.vertex[u] ^= 1;
stt.update_vertex(u);
memo.extend(g[u].iter().cloned());
}
}
} else {
dfs.push((v, save, true));
for (i, &u) in g[v].iter().enumerate() {
dfs.push((u, i == 0, false));
}
}
}
println!("{}", 2 * ans);
}
fn sack(v: usize, save: bool, g: &[Vec<usize>], stt: &mut STT<State>, ans: &mut u64) {
for (i, &u) in g[v].iter().enumerate().rev() {
sack(u, i == 0, g, stt, ans);
}
for &u in g[v].iter().skip(1) {
flip(u, g, stt);
}
stt.state.vertex[v] ^= 1;
stt.update_vertex(v);
*ans += stt.find().ans;
if !save {
flip(v, g, stt);
}
}
fn flip(v: usize, g: &[Vec<usize>], stt: &mut STT<State>) {
stt.state.vertex[v] ^= 1;
stt.update_vertex(v);
for &u in g[v].iter() {
flip(u, g, stt);
}
}
fn hld(g: &mut Vec<Vec<usize>>) {
let mut size = vec![1; g.len()];
let mut topo = vec![0];
for i in 0..g.len() {
let v = topo[i];
for u in g[v].clone() {
g[u].retain(|p| *p != v);
topo.push(u);
}
}
for &v in topo.iter().rev() {
g[v].sort_by_key(|u| !size[*u]);
for &u in g[v].iter() {
size[v] += size[u];
}
}
}
struct State {
vertex: Vec<u8>,
}
#[derive(Clone, Default, Debug)]
struct Data {
cnt: [u32; 2],
ls: [u64; 2],
rs: [u64; 2],
ans: u64,
len: u32,
}
impl TreeDP for State {
type Path = Data;
fn rake(&self, x: &Self::Path, y: &Self::Path) -> Self::Path {
let mut res = Data::default();
res.ans = x.ans + y.ans;
res.len = x.len;
for i in 0..2 {
res.cnt[i] = x.cnt[i] + y.cnt[i];
res.ls[i] = x.ls[i] + y.ls[i];
res.rs[i] = x.rs[i] + y.ls[i] + y.cnt[i] as u64 * x.len as u64;
res.ans += x.cnt[i] as u64 * y.ls[i ^ 1];
res.ans += x.ls[i] * y.cnt[i ^ 1] as u64;
}
res
}
fn compress(&self, p: &Self::Path, c: &Self::Path) -> Self::Path {
let mut res = Data::default();
res.ans = p.ans + c.ans;
res.len = p.len + c.len;
for i in 0..2 {
res.cnt[i] = p.cnt[i] + c.cnt[i];
res.ls[i] = p.ls[i] + c.ls[i] + p.len as u64 * c.cnt[i] as u64;
res.rs[i] = c.rs[i] + p.rs[i] + c.len as u64 * p.cnt[i] as u64;
res.ans += p.cnt[i] as u64 * c.ls[i ^ 1];
res.ans += p.rs[i] * c.cnt[i ^ 1] as u64;
}
res
}
fn build(&self, v: usize, _e: usize) -> Self::Path {
let x = self.vertex[v] as usize;
let mut res = Data::default();
res.cnt[x] += 1;
res.ls[x] += 1;
res.len = 1;
res
}
}
// ---------- 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 ----------
pub trait TreeDP {
type Path: Clone;
fn rake(&self, x: &Self::Path, y: &Self::Path) -> Self::Path;
fn compress(&self, p: &Self::Path, c: &Self::Path) -> Self::Path;
fn build(&self, v: usize, e: usize) -> Self::Path;
}
pub struct STT<R: TreeDP> {
node: UVec<Node<R::Path>>,
label: UVec<Label>,
parent: UVec<(usize, usize)>,
inv_edge: Vec<usize>,
state: R,
}
impl<R: TreeDP> STT<R> {
pub fn new(state: R, root: usize, edge: Vec<(usize, usize)>) -> Self {
let n = edge.len() + 1;
let mut deg = vec![0; n];
for &(a, b) in edge.iter() {
deg[a] += 1;
deg[b] += 1;
}
let mut g = deg
.into_iter()
.map(|d| Vec::with_capacity(d))
.collect::<Vec<_>>();
for (i, &(a, b)) in edge.iter().enumerate() {
g[a].push((b, i));
g[b].push((a, i));
}
let mut topo = vec![root];
let mut parent = vec![(n, !0); n];
let mut inv_edge = vec![!0; n - 1];
for i in 0..n {
let v = topo[i];
for i in 0..g[v].len() {
let (u, k) = g[v][i];
g[u].retain(|p| p.0 != v);
parent[u] = (v, k);
inv_edge[k] = u;
topo.push(u);
}
}
let mut size = vec![1; n];
for &v in topo.iter().rev() {
let c = &mut g[v];
let mut key = 0;
for i in 0..c.len() {
let (u, _) = c[i];
size[v] += size[u];
if key < size[u] {
key = size[u];
c.swap(0, i);
}
}
}
let init = Node::new(state.build(0, !0));
let mut node = vec![init.clone(); n];
let mut label = vec![Label::Vertex; n];
let mut merge = |x: usize, y: usize, op: Label| -> usize {
let k = node.len();
node.push(init.clone());
label.push(op);
node[x].p.set(k);
node[y].p.set(k);
node[k].l.set(x);
node[k].r.set(y);
k
};
let mut index = (0..n).collect::<Vec<_>>();
let mut depth = vec![0; n];
for &v in topo.iter().rev() {
let mut array: [Option<usize>; 128] = [None; 128];
let c = &g[v];
let mut up = 0usize;
let mut d = 0;
for &(u, _) in c.iter().rev() {
d = depth[u];
let mut pos = index[u];
while let Some(x) = array[d].take() {
pos = merge(pos, x, Label::Rake);
d += 1;
}
array[d] = Some(pos);
up = d.max(up);
}
let mut elem = false;
let mut r = None;
let mut res = 0;
for (i, a) in array[..=up].iter_mut().enumerate() {
if let Some(l) = a.take() {
if r.is_none() {
r = Some(l);
res = i;
elem = i == d;
continue;
}
if elem {
r = Some(merge(r.unwrap(), l, Label::Rake));
} else {
r = Some(merge(l, r.unwrap(), Label::Rake));
}
elem |= i == d;
res = i + 1;
}
}
if let Some(r) = r {
let u = c[0].0;
index[u] = r;
depth[u] = res;
}
if v == root || g[parent[v].0][0].0 != v {
let mut stack = vec![(depth[v], index[v])];
let mut pos = v;
while let Some(&(u, _)) = g[pos].get(0) {
pos = u;
stack.push((depth[u], index[u]));
while stack.len() >= 2 {
let k = stack.len();
if k >= 3
&& (stack[k - 3].0 == stack[k - 2].0
|| stack[k - 3].0 <= stack[k - 1].0)
{
let r = stack.pop().unwrap();
let m = stack.pop().unwrap();
let l = stack.pop().unwrap();
stack.push((m.0 + 1, merge(l.1, m.1, Label::Compress)));
stack.push(r);
} else if stack[k - 2].0 <= stack[k - 1].0 {
let m = stack.pop().unwrap();
let l = stack.pop().unwrap();
stack.push((m.0 + 1, merge(l.1, m.1, Label::Compress)));
} else {
break;
}
}
}
while stack.len() > 1 {
let r = stack.pop().unwrap();
let m = stack.pop().unwrap();
stack.push((m.0 + 1, merge(m.1, r.1, Label::Compress)));
}
let (a, b) = stack.pop().unwrap();
depth[v] = a;
index[v] = b;
}
}
for i in 0..n {
node[i].val = state.build(i, parent[i].1);
}
let mut res = Self {
node: UVec::new(node),
label: UVec::new(label),
state,
parent: UVec::new(parent),
inv_edge,
};
for i in n..res.node.len() {
res.pull(i);
}
res
}
pub fn update_vertex(&mut self, mut v: usize) {
self.node[v].val = self.state.build(v, self.parent[v].1);
while let Some(p) = self.node[v].p.get() {
self.pull(p);
v = p;
}
}
pub fn update_edge(&mut self, e: usize) {
self.update_vertex(self.inv_edge[e]);
}
pub fn find(&self) -> R::Path {
self.node.last().unwrap().val.clone()
}
fn pull(&mut self, v: usize) {
let l = self.node[v].l.get().unwrap();
let r = self.node[v].r.get().unwrap();
match self.label[v] {
Label::Rake => {
self.node[v].val = self.state.rake(&self.node[l].val, &self.node[r].val);
}
Label::Compress => {
self.node[v].val = self.state.compress(&self.node[l].val, &self.node[r].val);
}
_ => unreachable!(),
}
}
}
#[derive(Clone, Copy, Eq, PartialEq, Debug)]
enum Label {
Vertex,
Rake,
Compress,
}
#[derive(Clone, Debug)]
struct Node<T> {
val: T,
l: Pointer,
r: Pointer,
p: Pointer,
}
impl<T> Node<T> {
fn new(val: T) -> Self {
Self {
val,
l: Pointer::null(),
r: Pointer::null(),
p: Pointer::null(),
}
}
}
// ---------- begin pointer ----------
#[derive(Clone, Copy)]
pub struct Pointer(u32);
impl Pointer {
pub fn new(v: usize) -> Self {
Self(v as u32)
}
pub fn null() -> Self {
Self(!0)
}
pub fn get(&self) -> Option<usize> {
if self.0 == !0 {
None
} else {
Some(self.0 as usize)
}
}
pub fn is_null(&self) -> bool {
self.get().is_none()
}
pub fn set(&mut self, v: usize) {
self.0 = v as u32
}
}
impl From<usize> for Pointer {
fn from(x: usize) -> Self {
Self::new(x)
}
}
impl Default for Pointer {
fn default() -> Self {
Self::null()
}
}
impl std::fmt::Debug for Pointer {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
if let Some(x) = self.get() {
write!(f, "{}", x)
} else {
write!(f, "null")
}
}
}
// ---------- end pointer ----------
// reference: https://twitter.com/shino_skycrew/status/1322841105130422273
use std::ops::*;
#[derive(Clone, Eq, PartialEq, Ord, PartialOrd, Debug, Hash)]
pub struct UVec<T>(pub Vec<T>);
impl<T: Clone> UVec<T> {
pub fn new(ini: Vec<T>) -> Self {
UVec(ini)
}
}
impl<T> Deref for UVec<T> {
type Target = Vec<T>;
fn deref(&self) -> &Self::Target {
&self.0
}
}
impl<T> DerefMut for UVec<T> {
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.0
}
}
macro_rules! impl_index {
($x: ty) => {
impl<T> Index<$x> for UVec<T> {
type Output = T;
fn index(&self, index: $x) -> &Self::Output {
unsafe { self.0.get_unchecked(index as usize) }
}
}
impl<T> IndexMut<$x> for UVec<T> {
fn index_mut(&mut self, index: $x) -> &mut Self::Output {
unsafe { self.0.get_unchecked_mut(index as usize) }
}
}
};
}
impl_index!(usize);
impl_index!(u64);
impl_index!(u32);
impl_index!(u16);
impl_index!(u8);