AtCoder Beginner Contest 282 D

https://atcoder.jp/contests/abc282/tasks/abc282_d

連結成分に分けて、隣り合うノードを白と黒に塗り分けます。塗り分けられなかったら二部グラフではありません。そして、連結成分の中と連結成分の間にエッジをいくつ作れるかを数えます。

// Make Bipartite 2
#![allow(non_snake_case)]

use std::collections::HashMap;

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 read_input() -> (usize, Vec<(usize, usize)>) {
    let edge = |_| -> (usize, usize) {
        let v: Vec<usize> = read_vec();
        (v[0] - 1, v[1] - 1)
    };
    let v: Vec<usize> = read_vec();
    let N = v[0];
    let M = v[1];
    let edges = (0..M).map(edge).collect::<Vec<(usize, usize)>>();
    (N, edges)
}

type Graph = Vec<Vec<usize>>;

fn make_graph(N: usize, edges: Vec<(usize, usize)>) -> Graph {
    let mut graph: Graph = (0..N).map(|_| vec![]).collect();
    for (u, v) in edges.iter() {
        graph[*u].push(*v);
        graph[*v].push(*u)
    }
    graph
}

#[derive(PartialEq)]
enum Color {
    White, Black
}

fn reverse_color(color: &Color) -> Color {
    match *color {
        Color::White => Color::Black,
        Color::Black => Color::White,
    }
}

fn divide_into_connected(graph: &Graph) -> Vec<HashMap<usize, Color>> {
    let N = graph.len();
    let mut maps = Vec::<HashMap<usize, Color>>::new();
    let mut visited = (0..N).map(|_| false).collect::<Vec<bool>>();
    for v0 in 0..N {
        if visited[v0] {
            continue
        }
        
        let mut m = HashMap::<usize, Color>::new();
        m.insert(v0, Color::White);
        let mut stack: Vec<(usize, Color)> = vec![(v0, Color::White)];
        loop {
            match stack.pop() {
                Some((v, color)) => {
                    for v1 in graph.get(v).unwrap().iter() {
                        match m.get(&v1) {
                            Some(color1) => {
                                if *color1 == color {
                                    // 2部グラフでない
                                    return Vec::<HashMap<usize, Color>>::new();
                                }
                            },
                            None => {
                                stack.push((*v1, reverse_color(&color)));
                                m.insert(*v1, reverse_color(&color));
                                visited[*v1] = true
                            }
                        }
                    }
                },
                None => break
            }
        }
        maps.push(m)
    }
    maps
}

fn count_inner_connected(m: &HashMap<usize, Color>) -> usize {
    let num_white = m.iter().filter(|(_, c)| **c == Color::White).count();
    let num_black = m.iter().filter(|(_, c)| **c == Color::Black).count();
    num_white * num_black
}

fn count_inter_connected(maps: &Vec<HashMap<usize, Color>>) -> usize {
    let sum: usize = maps.iter().map(|m| m.len()).sum();
    let sum_squares: usize = maps.iter().map(|m| m.len() * m.len()).sum();
    (sum * sum - sum_squares) / 2
}

fn f(N: usize, edges: Vec<(usize, usize)>) -> usize {
    let M = edges.len();
    let graph = make_graph(N, edges);
    let maps = divide_into_connected(&graph);
    if maps.is_empty() {
        return 0
    }
    
    maps.iter().map(|m| count_inner_connected(m)).sum::<usize>() +
                            count_inter_connected(&maps) - M
}

fn main() {
    let v = read_input();
    println!("{}", f(v.0, v.1))
}