ABC223 G - Vertex Deletion

問題

atcoder.jp

解法

まず、木の最大マッチングは葉から貪欲に取っていくことで作れる。これは深さ優先探索で根付き木で (ある頂点vを含むサブツリーのマッチングの個数, vがマッチングに含まれているか) を再帰的に求めることで求まる。

この要領で、全方位木DPをして、全ての頂点について、その頂点を含まない(含むかどうか決める前の)マッチングの個数を求める。

コード

fn main() {
    let (r, w) = (std::io::stdin(), std::io::stdout());
    let mut sc = IO::new(r.lock(), w.lock());

    let n: usize = sc.read();
    let mut graph = vec![vec![]; n];
    for _ in 1..n {
        let a = sc.usize0();
        let b = sc.usize0();
        graph[a].push(b);
        graph[b].push(a);
    }

    let mut dp = vec![vec![]; n];
    let (greedy, _) = dfs(0, 0, &graph, &mut dp);

    let mut ans = vec![0; n];
    dfs2(0, 0, &graph, &mut dp, (0, false), &mut ans);

    let mut count = 0;
    for ans in ans {
        if ans == greedy {
            count += 1;
        }
    }
    println!("{}", count);
}

fn dfs(
    v: usize,
    p: usize,
    graph: &Vec<Vec<usize>>,
    dp: &mut Vec<Vec<(usize, bool)>>,
) -> (usize, bool) {
    let mut requested = false;
    let mut sum = 0;
    let deg = graph[v].len();
    dp[v] = vec![(0, false); deg];
    for (i, &next) in graph[v].iter().enumerate() {
        if next == p {
            continue;
        }

        let (count, prev_requested) = dfs(next, v, graph, dp);
        sum += count;
        requested |= prev_requested;
        dp[v][i] = (count, prev_requested);
    }

    if requested {
        (sum + 1, false)
    } else {
        (sum, true)
    }
}

fn dfs2(
    v: usize,
    p: usize,
    graph: &Vec<Vec<usize>>,
    dp: &mut Vec<Vec<(usize, bool)>>,
    dp_p: (usize, bool),
    ans: &mut [usize],
) {
    for (i, &next) in graph[v].iter().enumerate() {
        if next == p {
            dp[v][i] = dp_p;
        }
    }

    let deg = graph[v].len();
    let mut left_dp = vec![(0, false); deg + 1];
    let mut right_dp = vec![(0, false); deg + 1];

    for i in 0..deg {
        let (p_count, p_requested) = left_dp[i];
        let (dp_count, dp_requested) = dp[v][i];
        left_dp[i + 1] = (p_count + dp_count, p_requested || dp_requested);
    }
    for i in (0..deg).rev() {
        let (p_count, p_requested) = right_dp[i + 1];
        let (dp_count, dp_requested) = dp[v][i];
        right_dp[i] = (p_count + dp_count, p_requested || dp_requested);
    }

    ans[v] = left_dp[deg].0;
    for (i, &next) in graph[v].iter().enumerate() {
        if next == p {
            continue;
        }

        let (left_count, left_requested) = left_dp[i];
        let (right_count, right_requested) = right_dp[i + 1];

        let add = if left_requested || right_requested {
            1
        } else {
            0
        };

        dfs2(
            next,
            v,
            graph,
            dp,
            (left_count + right_count + add, add != 1),
            ans,
        );
    }
}

pub struct IO<R, W: std::io::Write>(R, std::io::BufWriter<W>);

impl<R: std::io::Read, W: std::io::Write> IO<R, W> {
    pub fn new(r: R, w: W) -> IO<R, W> {
        IO(r, std::io::BufWriter::new(w))
    }
    pub fn write<S: ToString>(&mut self, s: S) {
        use std::io::Write;
        self.1.write_all(s.to_string().as_bytes()).unwrap();
    }
    pub fn read<T: std::str::FromStr>(&mut self) -> T {
        use std::io::Read;
        let buf = self
            .0
            .by_ref()
            .bytes()
            .map(|b| b.unwrap())
            .skip_while(|&b| b == b' ' || b == b'\n' || b == b'\r' || b == b'\t')
            .take_while(|&b| b != b' ' && b != b'\n' && b != b'\r' && b != b'\t')
            .collect::<Vec<_>>();
        unsafe { std::str::from_utf8_unchecked(&buf) }
            .parse()
            .ok()
            .expect("Parse error.")
    }
    pub fn usize0(&mut self) -> usize {
        self.read::<usize>() - 1
    }
    pub fn vec<T: std::str::FromStr>(&mut self, n: usize) -> Vec<T> {
        (0..n).map(|_| self.read()).collect()
    }
    pub fn chars(&mut self) -> Vec<char> {
        self.read::<String>().chars().collect()
    }
}

コメント

全方位木DPは問題によって自由度が高いのでライブラリにしにくく感じる(一般にDPはライブラリにしにくい)ので、こうやって経験が積めるのは助かる。