use std::collections::HashMap;
fn main() {
let mut stones = vec![
vec![0, 0],
vec![0, 1],
vec![1, 0],
vec![1, 2],
vec![2, 1],
vec![2, 2],
];
let ans = remove_stones(&mut stones);
println!("ans = {}", ans);
}
fn remove_stones(stones: &mut Vec<Vec<i32>>) -> i32 {
let n = stones.len() as i32;
let mut row_pre: HashMap<i32, i32> = HashMap::new();
let mut col_pre: HashMap<i32, i32> = HashMap::new();
let mut uf = UnionFind::new(n);
for i in 0..n {
let x = stones[i as usize][0];
let y = stones[i as usize][1];
if !row_pre.contains_key(&x) {
row_pre.insert(x, i);
} else {
uf.union(i, *row_pre.get(&x).unwrap());
}
if !col_pre.contains_key(&y) {
col_pre.insert(y, i);
} else {
uf.union(i, *col_pre.get(&y).unwrap());
}
}
return n - uf.sets();
}
pub struct UnionFind {
father: Vec<i32>,
size: Vec<i32>,
help: Vec<i32>,
sets: i32,
}
impl UnionFind {
fn new(n: i32) -> UnionFind {
let mut father: Vec<i32> = vec![];
let mut size: Vec<i32> = vec![];
let mut help: Vec<i32> = vec![];
for i in 0..n {
father.push(i);
size.push(1);
help.push(0);
}
UnionFind {
father,
size,
help,
sets: n,
}
}
fn union(&mut self, mut i: i32, mut j: i32) {
i = self.find(i);
j = self.find(j);
if i != j {
if self.size[i as usize] >= self.size[j as usize] {
self.father[j as usize] = i;
self.size[i as usize] += self.size[j as usize];
} else {
self.father[i as usize] = j;
self.size[j as usize] += self.size[i as usize];
}
self.sets -= 1;
}
}
fn find(&mut self, mut i: i32) -> i32 {
let mut s = 0;
while i != self.father[i as usize] {
self.help[s] = i;
s += 1;
i = self.father[i as usize];
}
while s > 0 {
s -= 1;
self.father[self.help[s] as usize] = i;
}
return i;
}
fn sets(&mut self) -> i32 {
self.sets
}
}
评论