Union Find

今日練習的 1488. Avoid Flood in The City 可以用 Union Find 來增加一些執行效率，照例練習一下 C++/Python/Go/Rust 等程式語言。

C++

class UnionFind {
public:
    vector < int > root;
    UnionFind(int n) : root(n + 1) { iota(root.begin(), root.end(), 0); }
    int Find(int x) { return (x == root[x]) ? x : root[x] = Find(root[x]); }
    void UnionNext(int x) { root[x] = Find(x + 1); }
};
class Solution {
public:
    vector < int > avoidFlood(vector < int > & rains) {
        const int n = rains.size();
        UnionFind G(n);
        unordered_map < int, int > rainDay;
        rainDay.reserve(n);
        vector < int > ans(n, 1);
        for (int i = 0; i < n; i++) {
            const int lake = rains[i];
            if (lake > 0) {
                ans[i] = -1;
                G.UnionNext(i);
                auto it = rainDay.find(lake);
                if (it != rainDay.end()) {
                    int prev = it->second;
                    int dry = G.Find(prev + 1);
                    if (dry > i)
                        return {};
                    ans[dry] = lake;
                    G.UnionNext(dry);
                    it->second = i;
                } else {
                    rainDay[lake] = i;
                }
            }
        }
        return ans;
    }
};

Python3

class UnionFind:
    def __init__(self, n: int):
        self.root = [i for i in range(n + 1)]

    def find(self, x: int):
        if x == self.root[x]:
            return x
        else:
            self.root[x] = self.find(self.root[x])
            return self.root[x]

    def unionNext(self, x: int):
        self.root[x] = self.find(x + 1)


class Solution:
    def avoidFlood(self, rains: List[int]) -> List[int]:
        n = len(rains)
        G = UnionFind(n)
        rainDay = {}
        ans = [1] * n
        for i, lake in enumerate(rains):
            if lake > 0:
                ans[i] = -1
                G.unionNext(i)
                if lake in rainDay:
                    prev = rainDay[lake]
                    dry = G.find(prev + 1)
                    if dry > i:
                        return []
                    ans[dry] = lake
                    G.unionNext(dry)
                rainDay[lake] = i
        return ans

Go

type UnionFind []int

func (parent UnionFind) Find(x int) int {
	if x == parent[x] {
		return x
	} else {
		parent[x] = parent.Find(parent[x])
		return parent[x]
	}
}
func (parent UnionFind) UnionNext(x int) {
	parent[x] = parent.Find(x + 1)
}

func avoidFlood(rains []int) []int {
	n := len(rains)
	arr := make([]int, n+1)
	for i := range n + 1 {
		arr[i] = i
	}
	G := UnionFind(arr)
	rainDay := map[int]int{}
	ans := make([]int, n)
    for i := range n {
        ans[i] = 1
    }
	for i, lake := range rains {
		if lake > 0 {
			ans[i] = -1
			G.UnionNext(i)
			if prev, ok := rainDay[lake]; ok {
				dry := G.Find(prev + 1)
				if dry > i {
					return []int{}
				}
				ans[dry] = lake
				G.UnionNext(dry)
			}
			rainDay[lake] = i
		}
	}
	return ans
}

Rust

pub struct UnionFind {
    parent: Vec< i32 >,
}
impl UnionFind {
    pub fn new(n: usize) -> Self {
        let parent = (0..=n).map(|i| i as i32).collect();
        Self{parent}
    }
    pub fn find(&mut self, x: i32) -> i32 {
        let x_usize = x as usize;
        if self.parent[x_usize] == x {
            return x
        }
        self.parent[x_usize] = self.find(self.parent[x_usize]);
        self.parent[x_usize]
    }
    pub fn union_next(&mut self, x: i32) {
        self.parent[x as usize] = self.find(x + 1)
    }
}

use std::collections::HashMap;

impl Solution {
    pub fn avoid_flood(rains: Vec < i32 >) -> Vec < i32 > {
        let n = rains.len();
        let mut uf = UnionFind::new(n);
        let mut pool = HashMap::< i32,i32 >::new();
        let mut ans = vec![1; n];
        for (i, &lake) in rains.iter().enumerate() {
            if lake > 0 {
                ans[i] = -1;
                uf.union_next(i as i32);
                if let Some(prev) = pool.get(&lake) {
                    let dry = uf.find(prev + 1);
                    if dry > i as i32 {
                        return vec![];
                    }
                    ans[dry as usize] = lake;
                    uf.union_next(dry);
                }
                pool.insert(lake, i as i32);
            }
        }
        ans
    }
}