#724782 (Rust) No.1789 Tree Growing

• 問題: No.1789 Tree Growing
• ユーザー: akakimidori
• 提出日時: 2021-12-18 01:43:37
• 言語: Rust (1.83.0 + proconio)
• 結果: WA
• 実行時間: -
• コード長: 8,419 bytes
• コンパイル時間: 15,401 ms
• コンパイル使用メモリ: 377,052 KB
• 実行使用メモリ: 5,376 KB
• 最終ジャッジ日時: 2024-09-15 02:57:50
• 合計ジャッジ時間: 19,387 ms
• ジャッジサーバーID （参考情報）: judge3 / judge6

ファイルパターン	結果
sample	AC * 1 WA * 2
other	AC * 1 WA * 84

コンパイルメッセージ

warning: unused import: `std::io::Write`
   --> src/main.rs:177:5
    |
177 | use std::io::Write;
    |     ^^^^^^^^^^^^^^
    |
    = note: `#[warn(unused_imports)]` on by default

warning: type alias `Set` is never used
   --> src/main.rs:180:6
    |
180 | type Set<T> = BTreeSet<T>;
    |      ^^^
    |
    = note: `#[warn(dead_code)]` on by default

warning: type alias `Deque` is never used
   --> src/main.rs:181:6
    |
181 | type Deque<T> = VecDeque<T>;
    |      ^^^^^

ソースコード

// ---------- begin chmin, chmax ----------
pub trait ChangeMinMax {
    fn chmin(&mut self, x: Self) -> bool;
    fn chmax(&mut self, x: Self) -> bool;
}

impl<T: PartialOrd> ChangeMinMax for T {
    fn chmin(&mut self, x: Self) -> bool {
        *self > x && {
            *self = x;
            true
        }
    }
    fn chmax(&mut self, x: Self) -> bool {
        *self < x && {
            *self = x;
            true
        }
    }
}
// ---------- end chmin, chmax ----------
// ---------- begin recurse ----------
// reference
// https://twitter.com/noshi91/status/1393952665566994434
// https://twitter.com/shino16_cp/status/1393933468082397190
pub fn recurse<A, R, F>(f: F) -> impl Fn(A) -> R
where
    F: Fn(&dyn Fn(A) -> R, A) -> R,
{
    fn call<A, R, F>(f: &F, a: A) -> R
    where
        F: Fn(&dyn Fn(A) -> R, A) -> R,
    {
        f(&|a| call(f, a), a)
    }
    move |a| call(&f, a)
}
// ---------- end recurse ----------

// 最短路反復法
// src -> dst へflowだけ流せた時コストを返す
// 負閉路はないと仮定
// todo: SPFAじゃなくてポテンシャル付ダイクストラにする

const INF: i64 = 1_000_000i64 + 1;

struct Graph {
    size: usize,
    edge: Vec<(usize, usize, i64, i64)>,
}

impl Graph {
    fn new(size: usize) -> Self {
        Graph {
            size: size,
            edge: vec![],
        }
    }
    fn add_edge(&mut self, src: usize, dst: usize, capa: i64, cost: i64) {
        assert!(src < self.size && dst < self.size && src != dst);
        self.edge.push((src, dst, capa, cost));
    }
    fn solve(&self, src: usize, dst: usize, flow: i64) -> Option<i64> {
        if src == dst {
            return Some(0);
        }
        let size = self.size;
        let edge = &self.edge;
        let mut deg = vec![0; size];
        for &(a, b, _, _) in edge.iter() {
            deg[a] += 1;
            deg[b] += 1;
        }
        let mut graph: Vec<_> = deg.into_iter().map(|d| Vec::with_capacity(d)).collect();
        for &(a, b, capa, cost) in edge.iter() {
            let x = graph[a].len();
            let y = graph[b].len();
            graph[a].push((b, capa, cost, y));
            graph[b].push((a, 0, -cost, x));
        }
        let mut ans = 0;
        let mut dp = Vec::with_capacity(size);
        let mut elem = Vec::with_capacity(size);
        let mut que = std::collections::VecDeque::new();
        for _ in 0..flow {
            dp.clear();
            dp.resize(size, (INF, src, 0)); // コスト、親、親からの番号
            dp[src] = (0, src, 0);
            elem.clear();
            elem.resize(size, false);
            elem[src] = true;
            que.push_back(src);
            while let Some(v) = que.pop_front() {
                elem[v] = false;
                let (c, _, _) = dp[v];
                for (i, &(u, capa, cost, _)) in graph[v].iter().enumerate() {
                    if capa == 0 {
                        continue;
                    }
                    let c = c + cost;
                    if c < dp[u].0 {
                        dp[u] = (c, v, i);
                        if !elem[u] {
                            elem[u] = true;
                            que.push_back(u);
                        }
                    }
                }
            }
            if dp[dst].0 == INF {
                return None;
            }
            ans += dp[dst].0;
            let mut pos = dst;
            while pos != src {
                let (_, parent, k) = dp[pos];
                let inv = graph[parent][k].3;
                graph[parent][k].1 -= 1;
                graph[pos][inv].1 += 1;
                pos = parent;
            }
        }
        Some(ans)
    }
}
// ---------- begin input macro ----------
// reference: https://qiita.com/tanakh/items/0ba42c7ca36cd29d0ac8
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_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_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 ----------

use std::collections::*;
use std::io::Write;

type Map<K, V> = BTreeMap<K, V>;
type Set<T> = BTreeSet<T>;
type Deque<T> = VecDeque<T>;

fn run() {
    input! {
        k: usize,
        a: [(usize1, usize1); k - 1],
        n: usize,
        b: [(usize1, usize1); n - 1],
    }
    let mut g = vec![vec![]; k];
    for (a, b) in a {
        g[a].push(b);
        g[b].push(a);
    }
    let g = g;
    let mut h = vec![vec![]; n];
    for (a, b) in b {
        h[a].push(b);
        h[b].push(a);
    }
    let h = h;
    let mut topo = vec![0];
    let mut parent = vec![n; n];
    for i in 0..n {
        let v = topo[i];
        for &u in h[v].iter() {
            if u != parent[v] {
                parent[u] = v;
                topo.push(u);
            }
        }
    }
    let mut size = vec![0; n];
    for &v in topo.iter().rev() {
        let mut s = 1;
        for &u in h[v].iter() {
            s += size[u];
        }
        size[v] = s;
    }
    let size = |v: usize, p: usize| -> i64 {
        assert!(v < n && p <= n);
        if p == n {
            return n as i64;
        }
        if parent[v] == p {
            size[v] as i64
        } else if parent[p] == v {
            (n - size[p]) as i64
        } else {
            unreachable!("{} {}", v, p)
        }
    };
    let inf = 101;
    let memo = std::cell::RefCell::new(Map::new());
    let calc = recurse(|rec, (a, b, c, d): (usize, usize, usize, usize)| -> i64 {
        if let Some(&v) = memo.borrow().get(&(a, b, c, d)) {
            return v;
        }
        let d1 = if b < k { g[a].len() - 1 } else { g[a].len() };
        let d2 = if d < n { h[c].len() - 1 } else { h[c].len() };
        let mut res = inf;
        if b < k {
            let s = size(c, d);
            for &u in h[c].iter().filter(|p| **p != d) {
                res.chmin(s - size(u, c) - 1 + rec((a, b, u, c)));
            }
        }
        if d1 <= d2 {
            if d1 == 0 {
                res.chmin(size(c, d) - 1);
            } else {
                let mut graph = Graph::new(d1 + d2 + 2);
                let src = d1 + d2;
                let dst = src + 1;
                let mut geta = 0;
                for (i, &x) in g[a].iter().filter(|p| **p != b).enumerate() {
                    graph.add_edge(src, i, 1, 0);
                    for (j, &y) in h[c].iter().filter(|p| **p != d).enumerate() {
                        if i == 0 {
                            let s = size(y, c) as i64;
                            geta += s;
                            graph.add_edge(d1 + j, dst, 1, -s);
                        }
                        graph.add_edge(i, d1 + j, 1, rec((x, a, y, c)));
                    }
                }
                if let Some(v) = graph.solve(src, dst, d1 as i64) {
                    res.chmin(v + geta);
                }
            }
        }
        if res != inf {
            println!("{} {} {} {}: {}", a + 1, b + 1, c + 1, d + 1, res);
        }
        memo.borrow_mut().insert((a, b, c, d), res);
        res
    });
    let mut ans = inf;
    for v in 0..n {
        ans.chmin(calc((0, k, v, n)));
    }
    if ans == inf {
        println!("-1");
    } else {
        println!("{}", n as i64 - 1 - ans);
    }
}

fn main() {
    run();
}