AtCoder Beginner Contest 448 D

https://atcoder.jp/contests/abc448/tasks/abc448_d

頂点1から辿っていくだけですね。その際にAを追加していって、すぐに2つ以上あるAがあるかすぐに分かるようなデータ構造を作っておきます。

// Integer-duplicated Path
#![allow(non_snake_case)]


//////////////////// library ////////////////////

fn read<T: std::str::FromStr>() -> T {
    let mut line = String::new();
    std::io::stdin().read_line(&mut line).ok();
    line.trim().parse().ok().unwrap()
}

fn read_vec<T: std::str::FromStr>() -> Vec<T> {
    read::<String>().split_whitespace()
            .map(|e| e.parse().ok().unwrap()).collect()
}

fn YesNo(b: bool) -> String {
    return if b { "Yes" } else { "No" }.to_string()
}


//////////////////// Lables ////////////////////

use std::collections::{HashSet, HashMap};

struct Lables {
    set: HashSet<i32>,
    nums: HashMap<i32, u32>
}

impl Lables {
    fn contains_two_or_move(&self) -> bool {
        !self.nums.is_empty()
    }
    
    fn add(&mut self, a: i32) {
        if self.set.contains(&a) {
            self.set.remove(&a);
            self.nums.insert(a, 2);
        }
        else if self.nums.contains_key(&a) {
            let e = self.nums.entry(a).or_insert(0);
            *e += 1
        }
        else {
            self.set.insert(a);
        }
    }
    
    fn remove(&mut self, a: i32) {
        if self.set.contains(&a) {
            self.set.remove(&a);
        }
        else if self.nums.contains_key(&a) {
            let e = self.nums.entry(a).or_insert(0);
            if *e > 2 {
                *e -= 1
            }
            else {
                self.nums.remove(&a);
                self.set.insert(a);
            }
        }
    }
    
    fn new() -> Lables {
        Lables { set: HashSet::new(), nums: HashMap::new() }
    }
}


//////////////////// Graph ////////////////////

type Node = usize;
type Edge = (Node, Node);

fn read_edge() -> Edge {
    let v: Vec<usize> = read_vec();
    let (U, V) = (v[0]-1, v[1]-1);
    (U, V)
}

struct Graph {
    A: Vec<i32>,
    g: Vec<Vec<Node>>
}

impl Graph {
    fn from_edges(edges: Vec<Edge>, A: Vec<i32>) -> Graph{
        let N = A.len();
        let mut g: Vec<Vec<Node>> = vec![vec![]; N];
        for (u, v) in edges {
            g[u].push(v);
            g[v].push(u)
        }
        Graph { A, g }
    }
}


//////////////////// process ////////////////////

fn read_input() -> (Vec<i32>, Vec<Edge>) {
    let N: usize = read();
    let A: Vec<i32> = read_vec();
    let edges: Vec<Edge> = (0..N-1).map(|_| read_edge()).collect();
    (A, edges)
}

fn walk(v0: Node, lables: &mut Lables, results: &mut Vec<i32>, graph: &Graph) {
    lables.add(graph.A[v0]);
    results[v0] = if lables.contains_two_or_move() { 2 } else { 1 };
    for &v in graph.g[v0].iter() {
        if results[v] != 0 {
            continue
        }
        walk(v, lables, results, graph)
    }
    lables.remove(graph.A[v0])
}

fn F(A: Vec<i32>, edges: Vec<Edge>) {
    let N = A.len();
    let graph = Graph::from_edges(edges, A);
    let mut lables = Lables::new();
    // 0: unvisited 1: false 2: true
    let mut results: Vec<i32> = vec![0; N];
    walk(0, &mut lables, &mut results, &graph);
    for n in results {
        println!("{}", YesNo(n == 2))
    }
}

fn main() {
    let (A, edges) = read_input();
    F(A, edges)
}