2022-08-02：小红拿到了一个大立方体，该大立方体由1*1*1的小方块拼成，初始每个小方块都是白色。 小红可以每次选择一个小方块染成红色， 每次小红可能选

use rand::Rng;
fn main() {
    let size: i32 = 10;
    let test_time: i32 = 5000;
    println!("测试开始");
    for _ in 0..test_time {
        let n = rand::thread_rng().gen_range(0, size) + 1;
        let m = rand::thread_rng().gen_range(0, size) + 1;
        let h = rand::thread_rng().gen_range(0, size) + 1;
        let mut ops = random_ops(n, m, h);
        let ans1 = blocks1(n, m, h, &mut ops);
        let ans2 = blocks2(n, m, h, &mut ops);
        if ans1 != ans2 {
            println!("ans1 = {:?}", ans1);
            println!("ans2 = {:?}", ans2);
            println!("出错了!");
            break;
        }
    }
    println!("测试结束");
}

// 暴力方法
// 时间复杂度(k * n * m * h);
fn blocks1(n: i32, m: i32, h: i32, ops: &mut Vec<Vec<i32>>) -> Vec<i32> {
    let mut k = ops.len() as i32;
    //int[][][] cube = new int[n][m][h];
    let mut cube: Vec<Vec<Vec<i32>>> = vec![];
    for i in 0..n {
        cube.push(vec![]);
        for j in 0..m {
            cube[i as usize].push(vec![]);
            for _ in 0..h {
                cube[i as usize][j as usize].push(0);
            }
        }
    }
    let mut value = 1;
    let mut ans: Vec<i32> = vec![];
    for _ in 0..k {
        ans.push(0);
    }
    for i in 0..k {
        cube[ops[i as usize][0] as usize][ops[i as usize][1] as usize]
            [ops[i as usize][2] as usize] = -1;
        for x in 0..n {
            for y in 0..m {
                for z in 0..h {
                    if cube[x as usize][y as usize][z as usize] != -1
                        && cube[x as usize][y as usize][z as usize] != value
                    {
                        ans[i as usize] += 1;
                        infect(&mut cube, x, y, z, value);
                    }
                }
            }
        }
        value += 1;
    }
    return ans;
}

fn infect(cube: &mut Vec<Vec<Vec<i32>>>, a: i32, b: i32, c: i32, change: i32) {
    if a < 0
        || a == cube.len() as i32
        || b < 0
        || b == cube[0].len() as i32
        || c < 0
        || c == cube[0][0].len() as i32
        || cube[a as usize][b as usize][c as usize] == -1
        || cube[a as usize][b as usize][c as usize] == change
    {
        return;
    }
    cube[a as usize][b as usize][c as usize] = change;
    infect(cube, a - 1, b, c, change);
    infect(cube, a + 1, b, c, change);
    infect(cube, a, b - 1, c, change);
    infect(cube, a, b + 1, c, change);
    infect(cube, a, b, c - 1, change);
    infect(cube, a, b, c + 1, change);
}

// 最优解
// O(k + n * m * h)
fn blocks2(n: i32, m: i32, h: i32, ops: &mut Vec<Vec<i32>>) -> Vec<i32> {
    let mut k = ops.len() as i32;
    //int[][][] red = new int[n][m][h];
    let mut red: Vec<Vec<Vec<i32>>> = vec![];
    for i in 0..n {
        red.push(vec![]);
        for j in 0..m {
            red[i as usize].push(vec![]);
            for _ in 0..h {
                red[i as usize][j as usize].push(0);
            }
        }
    }
    for op in ops.iter() {
        red[op[0] as usize][op[1] as usize][op[2] as usize] += 1;
    }
    let mut uf = UnionFind::new(n, m, h, &mut red);
    let mut ans: Vec<i32> = vec![];
    for _ in 0..k {
        ans.push(0);
    }
    let mut i = k - 1;
    while i >= 0 {
        ans[i as usize] = uf.sets;
        let mut x = ops[i as usize][0];
        let mut y = ops[i as usize][1];
        let mut z = ops[i as usize][2];
        red[x as usize][y as usize][z as usize] -= 1;
        if red[x as usize][y as usize][z as usize] == 0 {
            // x, y ,z 这个格子，变白，建立自己的小集合
            // 然后6个方向，集合该合并合并
            uf.finger(x, y, z);
        }
        i -= 1;
    }
    return ans;
}

pub struct UnionFind {
    n: i32,
    m: i32,
    h: i32,
    father: Vec<i32>,
    size: Vec<i32>,
    help: Vec<i32>,
    sets: i32,
}
impl UnionFind {
    fn new(a: i32, b: i32, c: i32, red: &mut Vec<Vec<Vec<i32>>>) -> UnionFind {
        let mut n = a;
        let mut m = b;
        let mut h = c;
        let mut len = n * m * h;
        let mut father: Vec<i32> = vec![];
        let mut size: Vec<i32> = vec![];
        let mut help: Vec<i32> = vec![];
        for _ in 0..len {
            father.push(0);
            size.push(0);
            help.push(0);
        }
        let mut ans = UnionFind {
            n,
            m,
            h,
            father,
            size,
            help,
            sets: 0,
        };
        for x in 0..n {
            for y in 0..m {
                for z in 0..h {
                    if red[x as usize][y as usize][z as usize] == 0 {
                        ans.finger(x, y, z);
                    }
                }
            }
        }
        return ans;
    }

    pub fn finger(&mut self, x: i32, y: i32, z: i32) {
        // x，y，z
        // 一维数值
        let mut i = self.index(x, y, z);
        self.father[i as usize] = i;
        self.size[i as usize] = 1;
        self.sets += 1;
        self.union(i, x - 1, y, z);
        self.union(i, x + 1, y, z);
        self.union(i, x, y - 1, z);
        self.union(i, x, y + 1, z);
        self.union(i, x, y, z - 1);
        self.union(i, x, y, z + 1);
    }

    fn index(&mut self, x: i32, y: i32, z: i32) -> i32 {
        return z * self.n * self.m + y * self.n + x;
    }

    fn union(&mut self, mut i: i32, x: i32, y: i32, z: i32) {
        if x < 0 || x == self.n || y < 0 || y == self.m || z < 0 || z == self.h {
            return;
        }
        let mut j = self.index(x, y, z);
        if self.size[j as usize] == 0 {
            return;
        }
        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 random_ops(n: i32, m: i32, h: i32) -> Vec<Vec<i32>> {
    let mut size = rand::thread_rng().gen_range(0, n * m * h) + 1;
    let mut ans: Vec<Vec<i32>> = vec![];
    for i in 0..size {
        ans.push(vec![]);
        for _ in 0..3 {
            ans[i as usize].push(0);
        }
    }
    for i in 0..size {
        ans[i as usize][0] = rand::thread_rng().gen_range(0, n);
        ans[i as usize][1] = rand::thread_rng().gen_range(0, m);
        ans[i as usize][2] = rand::thread_rng().gen_range(0, h);
    }
    return ans;
}