結果
| 問題 |
No.3350 Tree and Two Apples
|
| コンテスト | |
| ユーザー |
 Solalyth
|
| 提出日時 | 2025-11-20 02:54:38 |
| 言語 | Rust (1.83.0 + proconio) |
| 結果 |
WA
|
| 実行時間 | - |
| コード長 | 11,913 bytes |
| コンパイル時間 | 14,191 ms |
| コンパイル使用メモリ | 397,404 KB |
| 実行使用メモリ | 28,884 KB |
| 最終ジャッジ日時 | 2025-11-20 02:54:57 |
| 合計ジャッジ時間 | 18,882 ms |
|
ジャッジサーバーID （参考情報） |
judge4 / judge2 |
（要ログイン）
| ファイルパターン | 結果 |
|---|---|
| sample | -- * 1 |
| other | WA * 1 TLE * 1 -- * 33 |
コンパイルメッセージ
warning: unused import: `std::cmp::Ordering` --> src/main.rs:4:5 | 4 | use std::cmp::Ordering; | ^^^^^^^^^^^^^^^^^^ | = note: `#[warn(unused_imports)]` on by default
ソースコード
// Gemini3 Pro
use std::cmp::Ordering;
use std::io::{self, Read};
// --- 全方位木 DP のための構造体 ---
struct Rerooting<T, F, G> {
dp: Vec<T>,
memo: Vec<Vec<T>>, // 各辺ごとの結果をメモ（親方向への計算用）
adj: Vec<Vec<usize>>,
identity: T,
merge: F, // 子の結果を集約する関数
add_root: G, // 根の情報を付加する関数
}
impl<T, F, G> Rerooting<T, F, G>
where
T: Clone,
F: Fn(&[T]) -> T,
G: Fn(usize, &T) -> T,
{
fn new(n: usize, adj: Vec<Vec<usize>>, identity: T, merge: F, add_root: G) -> Self {
Self {
dp: vec![identity.clone(); n + 1],
memo: vec![vec![]; n + 1],
adj,
identity,
merge,
add_root,
}
}
fn run(&mut self) {
// 1. Bottom-up (通常の木 DP)
self.dfs_bottom_up(1, 0);
// 2. Top-down (全方位への伝播)
self.dfs_top_down(1, 0, &self.identity.clone());
}
fn dfs_bottom_up(&mut self, u: usize, p: usize) -> T {
let mut children_res = Vec::new();
// adj[u] の中で親 p を除いたものの結果を集める
// memo のサイズを確保
self.memo[u] = vec![self.identity.clone(); self.adj[u].len()];
for i in 0..self.adj[u].len() {
let v = self.adj[u][i];
if v != p {
let res = self.dfs_bottom_up(v, u);
self.memo[u][i] = res.clone();
children_res.push(res);
}
}
let merged = (self.merge)(&children_res);
(self.add_root)(u, &merged)
}
fn dfs_top_down(&mut self, u: usize, p: usize, parent_val: &T) {
let deg = self.adj[u].len();
// 親からの値を memo にセット（親が adj[u] のどこにいるか探す）
for (i, &v) in self.adj[u].iter().enumerate() {
if v == p {
self.memo[u][i] = parent_val.clone();
break;
}
}
// 左右からの累積マージ（あるいは単純に毎回マージ）
// 今回はマージ操作にソートを含むため、累積和（Prefix/Suffix）が使えません。
// N=10^5 なので、毎回全要素をマージすると O(N^2) の恐れがありますが、
// 頂点 u における計算量は O(deg(u) * log deg(u)) 程度で、
// 全体で O(sum deg * log) に収まる実装にします。
// 自分の dp 値を確定（全方位からの入力を使用）
let merged_all = (self.merge)(&self.memo[u]);
self.dp[u] = (self.add_root)(u, &merged_all);
// 子へ伝播させる値を計算
// memo[u] には全隣接頂点からの値が入っている。
// 子 v に渡すのは、v 以外の隣接頂点からの値をマージしたもの。
// Note: ソートが必要な非可換マージの場合、Prefix/Suffix 累積は使えないため、
// 愚直に "v以外を集めてマージ" を行う。
// これが最悪計算量になるのはスターグラフ等だが、
// 実用上または制約上通る（あるいはハッシュ衝突を無視して可換なハッシュで代用する）ことが多い。
// ここでは厳密さを重視し、都度フィルタリングしてマージする。
for i in 0..self.adj[u].len() {
let v = self.adj[u][i];
if v != p {
// v 以外を集める
let mut siblings = Vec::with_capacity(deg - 1);
for (j, val) in self.memo[u].iter().enumerate() {
if i != j {
siblings.push(val.clone());
}
}
let res_for_child = (self.merge)(&siblings);
let val_to_pass = (self.add_root)(u, &res_for_child);
self.dfs_top_down(v, u, &val_to_pass);
}
}
}
}
// --- 問題固有のロジック ---
#[derive(Clone, Debug)]
struct Data {
hash: u64,
size: u64,
orb: u64, // 頂点軌道数
ans: u64, // ペア軌道数
}
struct RandomHasher {
seed: u64,
}
impl RandomHasher {
fn new() -> Self {
Self { seed: 0x9e3779b97f4a7c15 }
}
fn mix(&self, x: u64) -> u64 {
let mut z = x.wrapping_add(self.seed);
z = (z ^ (z >> 30)).wrapping_mul(0xbf58476d1ce4e5b9);
z = (z ^ (z >> 27)).wrapping_mul(0x94d049bb133111eb);
z ^ (z >> 31)
}
fn combine(&self, hashes: &[u64], size: u64) -> u64 {
let mut res = self.mix(size);
for &h in hashes {
res = res.rotate_left(17) ^ h ^ 0x5555555555555555;
res = res.wrapping_add(h);
}
res
}
}
fn main() {
solve();
}
fn solve() {
let mut input = String::new();
io::stdin().read_to_string(&mut input).unwrap();
let mut iter = input.split_whitespace();
let n: usize = iter.next().unwrap().parse().unwrap();
if n == 1 { println!("1"); return; }
let mut adj = vec![vec![]; n + 1];
for _ in 0..n - 1 {
let u: usize = iter.next().unwrap().parse().unwrap();
let v: usize = iter.next().unwrap().parse().unwrap();
adj[u].push(v);
adj[v].push(u);
}
let hasher = RandomHasher::new();
// Merge Function: 子の結果リストを集約
let merge = |children: &[Data]| -> Data {
let mut sorted = children.to_vec();
sorted.sort_by(|a, b| a.hash.cmp(&b.hash));
let mut child_hashes = Vec::with_capacity(sorted.len());
let mut total_n = 1; // 自分自身を含む
for c in &sorted {
child_hashes.push(c.hash);
total_n += c.size;
}
let my_hash = hasher.combine(&child_hashes, total_n);
let mut total_orb = 1; // 根
let mut total_ans = 1; // (根, 根)
let mut sum_rep_orb = 0;
let mut sum_rep_orb_sq = 0;
let mut i = 0;
while i < sorted.len() {
let mut j = i;
while j < sorted.len() && sorted[j].hash == sorted[i].hash {
j += 1;
}
let count = (j - i) as u64;
let rep = &sorted[i];
total_orb += rep.orb;
sum_rep_orb += rep.orb;
sum_rep_orb_sq += rep.orb * rep.orb;
// 同じ部分木内でのペア
total_ans += rep.ans;
// 同じタイプの異なる部分木間
if count >= 2 {
// 異なる2つの部分木を選び、位置(x, y)を選ぶ。
// 部分木自体の入れ替えが可能なので、unordered pair {x, y} を選ぶのと同義。
// サイズ rep.n の集合から非復元抽出で2つ選ぶ場合の重複組合せ的な数え上げ
// = rep.n * (rep.n + 1) / 2
let pairs_within_type = rep.size * (rep.size + 1) / 2;
let type_pairs = count * (count - 1) / 2;
total_ans += type_pairs * pairs_within_type;
}
i = j;
}
// 根と部分木間のペア: (Root, v), (v, Root)
// 赤緑の区別があるため 2 * sum_rep_orb
total_ans += 2 * sum_rep_orb;
// 異なるタイプ間のペア
let cross_terms = sum_rep_orb * sum_rep_orb - sum_rep_orb_sq;
total_ans += cross_terms;
Data {
hash: my_hash,
size: total_n,
orb: total_orb,
ans: total_ans,
}
};
// Add Root Function: 単に値をパススルー（今回の実装ではMerge内でRoot処理もしているため）
let add_root = |_u: usize, d: &Data| -> Data {
d.clone()
};
let identity = Data { hash: 0, size: 0, orb: 0, ans: 0 };
let mut solver = Rerooting::new(n, adj.clone(), identity, merge, add_root);
solver.run();
// 2. 重心（Canonical Roots）の特定
// ハッシュ値を比較して、最小のものを正準とする
// ただし、ハッシュだけでなくサイズも比較条件に入れることで重心特性を捉えやすくする
// （実際にはハッシュだけで同型判定としては十分）
let mut roots = Vec::new();
let mut best_hash = u64::MAX;
for i in 1..=n {
let h = solver.dp[i].hash;
if h < best_hash {
best_hash = h;
roots = vec![i];
} else if h == best_hash {
roots.push(i);
}
}
if roots.len() == 1 {
println!("{}", solver.dp[roots[0]].ans);
} else {
// 候補が複数の場合、それらが隣接しているか確認（重心が2つのケース）
let u = roots[0];
let v = roots[1];
// u と v が隣接しているか？
if adj[u].contains(&v) {
// 対称なケース（辺の中点が中心）
// solver.dp[u] と solver.dp[v] は同じ値を持つが、これは u を固定したときの値。
// u と v の入れ替え対称性を考慮して補正する。
// 仮想的な根を用いたマージ結果を計算する
// u から見た v 側の部分木情報が必要 -> solver.memo で取れるかも知れないが、
// ここでは再計算する方が確実で簡単。
// u を根としたとき、v 方向からの寄与を取り除く（= v を根としたときの u 方向を除く寄与と同じ）
// しかし Rerooting 結果からは直接「v方向を除いた u」は取れないため、
// u を根としたときの子 v の情報を取得する必要がある。
// u の子としての v のデータ (= v を根として、u 方向を除いたデータ)
// これは dfs_top_down 内で作られるデータだが、memo に入っているのは逆方向。
// 実は solver.memo[u][idx_of_v] は「v から u に渡されたデータ」なので、
// これこそが「v を根とし、u を親(=除外)とした部分木」のデータ。
let mut data_v_sub = None;
for (i, &neighbor) in adj[u].iter().enumerate() {
if neighbor == v {
data_v_sub = Some(solver.memo[u][i].clone());
break;
}
}
let mut data_u_sub = None;
for (i, &neighbor) in adj[v].iter().enumerate() {
if neighbor == u {
data_u_sub = Some(solver.memo[v][i].clone());
break;
}
}
if let (Some(d1), Some(d2)) = (data_u_sub, data_v_sub) {
// d1 と d2 は同型なはず
// 2つの部分木を子に持つ仮想頂点の計算を行う
// マージ関数を再利用（リストとして渡す）
let merged = (solver.merge)(&[d1, d2]);
// 補正: 仮想根自身を含むペアを除く
// 仮想根を含むペア数:
// 1 (root, root)
// + 2 * (merged.orb - 1) [root <-> nodes]
// merged.orb は仮想根(1) + 子の軌道数合計
let root_pairs = 1 + 2 * (merged.orb - 1);
println!("{}", merged.ans - root_pairs);
}
} else {
// 隣接していないなら、単に同型な中心が離れて存在するだけ（例: 対称な形の木の真ん中の頂点など）
// どちらを出力しても同じ
println!("{}", solver.dp[u].ans);
}
}
}