Submission #2508991

---

Source Code Expand

use std::io::*;
use std::str::FromStr;
use std::collections::HashMap;

#[allow(dead_code)]
fn read<T: FromStr>() -> T {
    let stdin = stdin();
    let stdin = stdin.lock();
    let token: String = stdin
        .bytes()
        .map(|c| c.expect("failed to read char") as char)
        .skip_while(|c| c.is_whitespace())
        .take_while(|c| !c.is_whitespace())
        .collect();
    token.parse().ok().expect("failed to parse token")
}

type Weight = i32;
#[allow(dead_code)]
#[derive(Clone, Debug, Eq, PartialEq)]
struct Edge {
    index: usize,
    src: usize,
    dest: usize,
    weight: Weight,
}
#[allow(dead_code)]
impl Edge {
    fn new(index: usize, src: usize, dest: usize, weight: Weight) -> Edge {
        Edge {
            index: index,
            src: src,
            dest: dest,
            weight: weight,
        }
    }
}
impl std::cmp::Ord for Edge {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        (self.weight, self.src, self.dest).cmp(&(other.weight, other.src, other.dest))
    }
}
impl std::cmp::PartialOrd for Edge {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}
type Edges = Vec<Edge>;
type Graph = Vec<Edges>;
#[allow(dead_code)]
fn graph_new(n: usize) -> Graph {
    vec![vec![]; n]
}
#[allow(dead_code)]
fn add_bi_edge(graph: &mut Graph, index: &mut usize, src: usize, dest: usize, weight: Weight) {
    graph[src].push(Edge::new(*index, src, dest, weight));
    *index += 1;
    graph[dest].push(Edge::new(*index, dest, src, weight));
    *index += 1;
}

struct Solver {
    n: usize,
    g: Graph,
    edge_index: HashMap<(usize, usize), usize>,
    initial_colors: Vec<usize>,
    memo: Vec<Vec<i64>>,
    memo_all_black: Vec<Option<bool>>,
    memo_total: Vec<(usize, i64)>,
    memo_childs: Vec<Vec<(i64, i64, i64, usize)>>,
    memo_all_black_total: Vec<i64>,
    visit_cnt: Vec<Vec<i64>>,
    visit_cnt_all_black: Vec<i64>,
}

impl Solver {
    fn new() -> Solver {
        let n = read::<usize>();
        let mut g = graph_new(n);
        let mut edge_index = HashMap::<(usize, usize), usize>::new();
        let mut e = 0;
        for _ in 0..n - 1 {
            let x = read::<usize>() - 1;
            let y = read::<usize>() - 1;
            edge_index.insert((x, y), e);
            edge_index.insert((y, x), e + 1);
            add_bi_edge(&mut g, &mut e, x, y, 0);
        }
        let initial_colors = read::<String>()
            .chars()
            .map(|c| if c == 'W' { 0 } else { 1 })
            .collect::<Vec<usize>>();
        Solver {
            n: n,
            g: g,
            edge_index: edge_index,
            initial_colors: initial_colors,
            memo: vec![vec![-1; 2 * (n - 1)]; 2],
            memo_all_black: vec![None; 2 * (n - 1)],
            memo_total: vec![(0, -1); n],
            memo_childs: vec![vec![]; n],
            memo_all_black_total: vec![-1; n],
            visit_cnt: vec![vec![0; n]; 2],
            visit_cnt_all_black: vec![0; n],
        }
    }
    fn move_edge_all_black(&mut self, edge: &Edge) -> bool {
        if self.memo_all_black[edge.index].is_some() {
            return self.memo_all_black[edge.index].unwrap();
        }
        let ret = self.all_black(edge.dest, edge.src);
        self.memo_all_black[edge.index] = Some(ret);
        ret
    }
    fn all_black(&mut self, from: usize, parent: usize) -> bool {
        if self.initial_colors[from] == 0 {
            return false;
        }
        self.visit_cnt_all_black[from] += 1;
        if self.visit_cnt_all_black[from] > 2 {
            if self.memo_all_black_total[from] == -1 {
                let mut sum = 0;
                for edge in self.g[from].clone().iter() {
                    sum += if self.move_edge_all_black(edge) { 1 } else { 0 };
                }
                self.memo_all_black_total[from] = sum;
            }
            let rev_index = self.edge_index[&(from, parent)];
            let mut cnt = self.memo_all_black_total[from];
            cnt -= if self.memo_all_black[rev_index].unwrap() {
                1
            } else {
                0
            };
            return cnt == self.g[from].len() as i64 - 1;
        }
        for edge in self.g[from].clone().iter() {
            if edge.dest == parent {
                continue;
            }
            if !self.move_edge_all_black(edge) {
                return false;
            }
        }
        true
    }
    fn move_edge(&mut self, back: usize, edge: &Edge) -> i64 {
        if self.memo[back][edge.index] != -1 {
            return self.memo[back][edge.index];
        }
        let c = self.initial_colors[edge.dest] ^ 1;
        let ret = back as i64 + 1 + self.dfs(back, c, edge.dest, edge.src);
        self.memo[back][edge.index] = ret;
        ret
    }
    fn dfs(&mut self, back: usize, c: usize, from: usize, parent: usize) -> i64 {
        self.visit_cnt[back][from] += 1;
        if c != self.initial_colors[from] && self.visit_cnt[back][from] > 2 {
            let mut c = c;
            let mut childs = vec![];
            if self.memo_total[from].1 == -1 {
                let mut sum = 0;
                for edge in self.g[from].clone().iter() {
                    if self.move_edge_all_black(edge) {
                        continue;
                    }
                    c ^= 1;
                    let one_way = self.move_edge(0, edge);
                    let need_back = self.move_edge(1, edge);
                    let diff = need_back - one_way;
                    childs.push((diff, need_back, one_way, edge.dest));
                    sum += need_back;
                }
                childs.sort();
                childs.reverse();
                self.memo_total[from] = (c, sum);
                self.memo_childs[from] = childs;
            }
            let rev_index = self.edge_index[&(from, parent)];
            let mut c = self.memo_total[from].0;
            let mut ret = self.memo_total[from].1;
            if !self.memo_all_black[rev_index].unwrap() {
                c ^= 1;
                ret -= self.memo[1][rev_index];
            }
            if back == 0 && self.memo_childs[from].len() > 0 {
                if self.memo_childs[from][0].3 != parent {
                    c ^= 1;
                    ret -= self.memo_childs[from][0].0;
                } else if self.memo_childs[from].len() > 1 {
                    c ^= 1;
                    ret -= self.memo_childs[from][1].0;
                }
            }
            if c == 0 {
                ret += 1;
            }
            return ret;
        }
        let mut c = c;
        let mut ret = 0;
        let mut childs = vec![];
        for edge in self.g[from].clone().iter() {
            if edge.dest == parent || self.move_edge_all_black(edge) {
                continue;
            }
            c ^= 1;
            let one_way = self.move_edge(0, edge);
            let need_back = self.move_edge(1, edge);
            let diff = need_back - one_way;
            childs.push((diff, need_back, one_way));
            ret += need_back;
        }
        childs.sort();
        childs.reverse();
        if back == 0 && childs.len() > 0 {
            // println!("{} {:?}", from, childs);
            c ^= 1;
            ret -= childs[0].0;
        }
        if c == 0 {
            ret += 1;
        }
        ret
    }
    fn solve(&mut self) {
        let mut ans = 1 << 60;
        for i in 0..self.n {
            let c = self.initial_colors[i];
            let nret = self.dfs(0, c, i, i);
            ans = std::cmp::min(ans, nret);
        }
        println!("{}", ans);
    }
}

fn main() {
    let mut solver = Solver::new();
    solver.solve();
}

Submission Info

Judge Result