競プロ典型 012

https://atcoder.jp/contests/typical90/tasks/typical90_l

典型的なUnion-Findを使う問題ですね。
色々な型に対応できるようにUnion-Findの構造体を頑張って書いてみました。

// Gravy Jobs
#![allow(non_snake_case)]

use std::collections::HashMap;
use std::hash::Hash;


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


//////////////////// Union-Find ////////////////////

struct UnionFind<T: Copy + Eq + Hash> {
    parents: HashMap<T, T>,
    heights: HashMap<T, i32>
}

impl<T: Copy + Eq + Hash> UnionFind<T> {
    fn join(&mut self, v1: T, v2: T) {
        let r1 = self.root(v1);
        let r2 = self.root(v2);
        let h1 = self.height(v1);
        let h2 = self.height(v2);
        if h1 == h2 {
            self.parents.insert(r1, r2);
            self.heights.insert(r1, h1 + 1);
        }
        else if h1 < h2 {
            self.parents.insert(r1, r2);
        }
        else {
            self.parents.insert(r2, r1);
        }
    }
    
    fn root(&self, mut v: T) -> T {
        loop {
            let &v1 = self.parents.get(&v).unwrap();
            if v1 == v {
                break
            }
            v = v1
        }
        v
    }
    
    fn height(&self, v: T) -> i32 {
        *self.heights.get(&v).unwrap()
    }
    
    fn new(nodes: Vec<T>) -> UnionFind<T> {
        let parents: HashMap<T, T> = nodes.iter().map(|&v| (v, v)).collect();
        let heights: HashMap<T, i32> = nodes.iter().map(|&v| (v, 1)).collect();
        UnionFind { parents, heights }
    }
}


//////////////////// Board ////////////////////

type Point = (usize, usize);

struct Board {
    H: usize,
    W: usize,
    colors: Vec<Vec<bool>>,
    uf: UnionFind<Point>
}

impl Board {
    fn new(H: usize, W: usize) -> Board {
        let colors: Vec<Vec<bool>> = (0..H).map(|_| vec![false; W]).collect();
        let nodes: Vec<Point>
                = (0..H).flat_map(|x| (0..W).map(move |y| (x, y))).collect();

        let uf: UnionFind<Point> = UnionFind::<Point>::new(nodes);
        Board { H, W, colors, uf }
    }
    
    fn query(&mut self, query: Query) {
        match query {
            Query::Paint(pt) => self.paint(pt),
            Query::Test(pt1, pt2) => self.test_connect(pt1, pt2)
        }
    }
    
    fn paint(&mut self, pt: Point) {
        self.colors[pt.0][pt.1] = true;
        for pt1 in self.neighbors(&pt) {
            if self.is_red(&pt1) {
                self.uf.join(pt, pt1)
            }
        }
    }
    
    fn test_connect(&self, pt1: Point, pt2: Point) {
        if !self.is_red(&pt1) || !self.is_red(&pt2) {
            println!("{}", YesNo(false))
        }
        else {
            let r1 = self.uf.root(pt1);
            let r2 = self.uf.root(pt2);
            println!("{}", YesNo(r1 == r2))
        }
    }
    
    fn neighbors(&self, pt: &Point) -> Vec<Point> {
        let mut pts: Vec<Point> = vec![];
        if pt.0 > 0 {
            pts.push((pt.0 - 1, pt.1))
        }
        if pt.0 < self.H - 1 {
            pts.push((pt.0 + 1, pt.1))
        }
        if pt.1 > 0 {
            pts.push((pt.0, pt.1 - 1))
        }
        if pt.1 < self.W - 1 {
            pts.push((pt.0, pt.1 + 1))
        }
        pts
    }
    
    fn is_red(&self, pt: &Point) -> bool {
        self.colors[pt.0][pt.1]
    }
}


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

enum Query {
    Paint(Point),
    Test(Point, Point),
}

impl Query {
    fn read() -> Query {
        let v: Vec<usize> = read_vec();
        if v[0] == 1 {
            Query::Paint((v[1]-1, v[2]-1))
        }
        else {
            Query::Test((v[1]-1, v[2]-1), (v[3]-1, v[4]-1))
        }
    }
}


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

fn read_input() -> (usize, usize, usize) {
    let v = read_vec();
    let (H, W) = (v[0], v[1]);
    let Q = read();
    (H, W, Q)
}

fn F(H: usize, W: usize, Q: usize) {
    let mut board = Board::new(H, W);
    for _ in 0..Q {
        let query = Query::read();
        board.query(query)
    }
}

fn main() {
    let (H, W, Q) = read_input();
    F(H, W, Q)
}