結果

問題 No.2892 Lime and Karin
ユーザー koba-e964koba-e964
提出日時 2024-10-22 17:35:41
言語 Rust
(1.83.0 + proconio)
結果
WA  
実行時間 -
コード長 6,207 bytes
コンパイル時間 15,163 ms
コンパイル使用メモリ 401,040 KB
実行使用メモリ 31,296 KB
最終ジャッジ日時 2024-10-22 17:36:09
合計ジャッジ時間 25,131 ms
ジャッジサーバーID
(参考情報)
judge3 / judge4
このコードへのチャレンジ
(要ログイン)
ファイルパターン 結果
sample AC * 3
other AC * 17 WA * 35
権限があれば一括ダウンロードができます
コンパイルメッセージ
warning: unused import: `BufWriter`
 --> src/main.rs:5:22
  |
5 | use std::io::{Write, BufWriter};
  |                      ^^^^^^^^^
  |
  = note: `#[warn(unused_imports)]` on by default

warning: unused import: `Write`
 --> src/main.rs:5:15
  |
5 | use std::io::{Write, BufWriter};
  |               ^^^^^

ソースコード

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

#[allow(unused_imports)]
use std::cmp::*;
#[allow(unused_imports)]
use std::collections::*;
use std::io::{Write, BufWriter};
// https://qiita.com/tanakh/items/0ba42c7ca36cd29d0ac8
macro_rules! input {
($($r:tt)*) => {
let stdin = std::io::stdin();
let mut bytes = std::io::Read::bytes(std::io::BufReader::new(stdin.lock()));
let mut next = move || -> String{
bytes.by_ref().map(|r|r.unwrap() as char)
.skip_while(|c|c.is_whitespace())
.take_while(|c|!c.is_whitespace())
.collect()
};
input_inner!{next, $($r)*}
};
}
macro_rules! input_inner {
($next:expr) => {};
($next:expr,) => {};
($next:expr, $var:ident : $t:tt $($r:tt)*) => {
let $var = read_value!($next, $t);
input_inner!{$next $($r)*}
};
}
macro_rules! read_value {
($next:expr, ( $($t:tt),* )) => { ($(read_value!($next, $t)),*) };
($next:expr, [ $t:tt ; $len:expr ]) => {
(0..$len).map(|_| read_value!($next, $t)).collect::<Vec<_>>()
};
($next:expr, chars) => {
read_value!($next, String).chars().collect::<Vec<char>>()
};
($next:expr, usize1) => (read_value!($next, usize) - 1);
($next:expr, $t:ty) => ($next().parse::<$t>().expect("Parse error"));
}
trait Bisect<T> {
fn lower_bound(&self, val: &T) -> usize;
fn upper_bound(&self, val: &T) -> usize;
}
impl<T: Ord> Bisect<T> for [T] {
fn lower_bound(&self, val: &T) -> usize {
let mut pass = self.len() + 1;
let mut fail = 0;
while pass - fail > 1 {
let mid = (pass + fail) / 2;
if &self[mid - 1] >= val {
pass = mid;
} else {
fail = mid;
}
}
pass - 1
}
fn upper_bound(&self, val: &T) -> usize {
let mut pass = self.len() + 1;
let mut fail = 0;
while pass - fail > 1 {
let mid = (pass + fail) / 2;
if &self[mid - 1] > val {
pass = mid;
} else {
fail = mid;
}
}
pass - 1
}
}
// Returns (root, children)
// This functions uses O(n) stack space.
// Complexity: O(n log n)-time, O(n)-space
// Verified by: ABC291-Ex (https://atcoder.jp/contests/abc291/submissions/39303290)
fn centroid_decompose(g: &[Vec<usize>]) -> (usize, Vec<Vec<usize>>) {
fn find_subtree_sizes(g: &[Vec<usize>], v: usize, par: usize,
dp: &mut [usize], vis: &[bool]) {
let mut sum = 1;
for &w in &g[v] {
if par == w || vis[w] { continue; }
find_subtree_sizes(g, w, v, dp, vis);
sum += dp[w];
}
dp[v] = sum;
}
fn centroid_decompose_inner(g: &[Vec<usize>], v: usize, par: usize,
ch: &mut [Vec<usize>],
dp: &mut [usize], vis: &mut [bool]) -> usize {
let n = g.len();
find_subtree_sizes(g, v, n, dp, vis);
let cent = {
let sz = dp[v];
let find_centroid = |mut v: usize, mut par: usize| {
loop {
let mut has_majority = false;
for &w in &g[v] {
if par == w || vis[w] { continue; }
if dp[w] > sz / 2 {
par = v;
v = w;
has_majority = true;
break;
}
}
if !has_majority {
return v;
}
}
};
find_centroid(v, n)
};
if par < n {
ch[par].push(cent);
}
// v was selected as a centroid
// and will be ignored in the following decomposition procedure
vis[cent] = true;
for &w in &g[cent] {
if !vis[w] {
centroid_decompose_inner(g, w, cent, ch, dp, vis);
}
}
cent
}
let n = g.len();
let mut ch = vec![vec![]; n];
// This Vec is used across multiple calls to `centroid_decompose_inner`
let mut dp = vec![0; n];
let mut vis = vec![false; n];
let root = centroid_decompose_inner(&g, 0, n, &mut ch, &mut dp, &mut vis);
(root, ch)
}
fn main() {
// In order to avoid potential stack overflow, spawn a new thread.
let stack_size = 104_857_600; // 100 MB
let thd = std::thread::Builder::new().stack_size(stack_size);
thd.spawn(|| solve()).unwrap().join().unwrap();
}
fn dfs(v: usize, par: usize, g: &[Vec<usize>], s: &[char], vis: &[bool], cur: i32, agg: &mut Vec<i32>) {
let cur = if s[v] == '1' { cur + 1 } else { cur - 1 };
agg.push(cur);
for &w in &g[v] {
if par == w || vis[w] { continue; }
dfs(w, v, g, s, vis, cur, agg);
}
}
fn calc_comb(init: i32, mut comb: Vec<i32>) -> i64 {
comb.sort_unstable();
let mut ret = 0;
for &v in &comb {
let idx = comb.lower_bound(&(-init - v + 1));
ret += (comb.len() - idx) as i64;
}
ret
}
// https://yukicoder.me/problems/no/2892 (4)
fn solve() {
input! {
n: usize,
uv: [(usize1, usize1); n - 1],
s: chars,
}
let mut g = vec![vec![]; n];
for &(u, v) in &uv {
g[u].push(v);
g[v].push(u);
}
let (root, ch) = centroid_decompose(&g);
let mut vis = vec![false; n];
let mut que = vec![root];
let mut tot = 0;
while let Some(v) = que.pop() {
vis[v] = true;
for &w in &ch[v] {
que.push(w);
}
// search
let mut agg = vec![];
let init = if s[v] == '1' { 1 } else { -1 };
for &c in &ch[v] {
let mut me = vec![];
dfs(c, v, &g, &s, &vis, 0, &mut me);
agg.extend_from_slice(&me);
tot -= calc_comb(init, me);
}
for &a in &agg {
if init + a > 0 {
tot += 2;
}
}
tot += calc_comb(init, agg);
if init == 1 {
tot += 2;
}
}
println!("{}", tot / 2);
}
הההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההה
XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
0