#724846 (Rust) No.1789 Tree Growing

• 問題: No.1789 Tree Growing
• ユーザー: akakimidori
• 提出日時: 2021-12-18 14:39:28
• 言語: Rust (1.83.0 + proconio)
• 結果: AC
• 実行時間: 565 ms / 5,000 ms
• コード長: 9,435 bytes
• コンパイル時間: 14,498 ms
• コンパイル使用メモリ: 378,308 KB
• 実行使用メモリ: 5,376 KB
• 最終ジャッジ日時: 2024-09-15 13:47:52
• 合計ジャッジ時間: 18,028 ms
• ジャッジサーバーID （参考情報）: judge3 / judge2

ファイルパターン	結果
sample	AC * 3
other	AC * 85

ソースコード

// ---------- 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 ----------

const INF: i64 = 1_000_000_000_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 flow(&self, src: usize, dst: usize, flow: i64) -> Vec<(i64, i64)> {
        assert!(src != dst);
        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 sum = 0;
        let mut res = vec![];
        let mut used = 0;
        let mut dp = Vec::with_capacity(size);
        let mut pot = vec![0; size];
        let mut h = std::collections::BinaryHeap::new();
        while used < flow {
            dp.clear();
            dp.resize(size, (INF, src, 0)); // コスト、親、親からの番号
            dp[src] = (0, src, 0);
            if used > 0 {
                h.push((0, src));
                while let Some((d, v)) = h.pop() {
                    let d = -d;
                    if d > dp[v].0 {
                        continue;
                    }
                    for (i, &(u, capa, cost, _)) in graph[v].iter().enumerate() {
                        if capa == 0 {
                            continue;
                        }
                        let c = d + cost - pot[u] + pot[v];
                        if c < dp[u].0 {
                            dp[u] = (c, v, i);
                            h.push((-c, u));
                        }
                    }
                }
            } else {
                let mut deg = vec![0; size];
                for &(dst, capa, _, _) in graph.iter().flat_map(|g| g.iter()) {
                    if capa > 0 {
                        deg[dst] += 1;
                    }
                }
                let mut dfs = vec![];
                for (i, deg) in deg.iter().enumerate() {
                    if *deg == 0 {
                        dfs.push(i);
                    }
                }
                while let Some(v) = dfs.pop() {
                    for (i, &(u, capa, cost, _)) in graph[v].iter().enumerate() {
                        if capa == 0 {
                            continue;
                        }
                        let c = dp[v].0 + cost;
                        if c < dp[u].0 {
                            dp[u] = (c, v, i);
                        }
                        deg[u] -= 1;
                        if deg[u] == 0 {
                            dfs.push(u);
                        }
                    }
                }
            }
            if dp[dst].0 == INF {
                break;
            }
            for i in 0..size {
                pot[i] -= dp[dst].0 - dp[i].0;
            }
            let mut sub = 1;
            let mut pos = dst;
            while pos != src {
                let (_, parent, k) = dp[pos];
                sub = std::cmp::min(sub, graph[parent][k].1);
                pos = parent;
            }
            let mut pos = dst;
            while pos != src {
                let (_, parent, k) = dp[pos];
                let inv = graph[parent][k].3;
                graph[parent][k].1 -= sub;
                graph[pos][inv].1 += sub;
                pos = parent;
            }
            used += sub;
            sum += -pot[src] * sub;
            res.push((used, sum));
        }
        res
    }
}
// ---------- 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 ----------

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(std::collections::HashMap::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(s) = graph.flow(src, dst, d1 as i64).last() {
                    res.chmin(s.1 + geta);
                }
            }
        }
        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();
}