AtCoder Beginner Contest 303 E

https://atcoder.jp/contests/abc303/tasks/abc303_e

最初にグラフを作るのではなく、エッジから直接次数が2より大きいノードを選んで、そのエッジを取り除いてグラフを作ります。

// A Gift From the Stars
#![allow(non_snake_case)]

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


//////////////////// 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()
}


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

type Node = usize;
type Edge = (Node, Node);
type Graph = HashMap<Node, Vec<Node>>;

fn make_graph(edges: Vec<Edge>) -> Graph {
    let mut graph = Graph::new();
    for (u, v) in edges.into_iter() {
        let e1 = graph.entry(u).or_insert(vec![]);
        (*e1).push(v);
        let e2 = graph.entry(v).or_insert(vec![]);
        (*e2).push(u);
    }
    graph
}

fn divide_into_connected(graph: &Graph) -> Vec<Graph> {
    let mut subgraphs = Vec::<Graph>::new();
    let mut unvisited = graph.keys().map(|v| *v).collect::<HashSet<Node>>();
    while !unvisited.is_empty() {
        let v0 = unvisited.iter().next().unwrap();
        let mut g = Graph::new();
        let mut stack: Vec<Node> = vec![*v0];
        while let Some(v) = stack.pop() {
            let vs = graph.get(&v).unwrap();
            let vs_copy = vs.iter().map(|v1| *v1).collect::<Vec<Node>>();
            for v1 in vs_copy.iter() {
                if unvisited.contains(v1) {
                    stack.push(*v1);
                    unvisited.remove(v1);
                }
            }
            g.insert(v, vs_copy);
        }
        subgraphs.push(g)
    }
    subgraphs
}


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

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

fn f(edges: Vec<Edge>) -> Vec<u32> {
    let mut counter: HashMap<Node, u32> = HashMap::new();
    for (u, v) in edges.iter() {
        let e1 = counter.entry(*u).or_insert(0);
        *e1 += 1;
        let e2 = counter.entry(*v).or_insert(0);
        *e2 += 1;
    }
    
    let mut ns: Vec<u32> = Vec::new();
    let mut vs: HashSet<Node> = HashSet::new();
    for (v, freq) in counter.into_iter() {
        if freq > 2 {
            ns.push(freq);
            vs.insert(v);
        }
    }
    
    let mut new_edges: Vec<Edge> = Vec::new();
    for (u, v) in edges.into_iter() {
        if !vs.contains(&u) && !vs.contains(&v) {
            new_edges.push((u, v))
        }
    }
    
    let graph = make_graph(new_edges);
    let subgraphs = divide_into_connected(&graph);
    for subgraph in subgraphs.into_iter() {
        for _ in 0..subgraph.len()/3 {
            ns.push(2)
        }
    }
    ns.sort();
    ns
}


//////////////////// main ////////////////////

fn main() {
    let edges = read_input();
    let ns = f(edges);
    let ss: Vec<String> = ns.iter().map(|n| n.to_string()).collect();
    println!("{}", ss.join(" "))
}