rustgym 0.2.0

rustgym solutions
Documentation
struct Solution;

use std::collections::BTreeMap;
use std::collections::BTreeSet;
use std::collections::HashMap;

struct UnionFind {
    parent: Vec<usize>,
    n: usize,
}

impl UnionFind {
    fn new(n: usize) -> Self {
        let parent = (0..n).collect();
        UnionFind { parent, n }
    }

    fn find(&mut self, i: usize) -> usize {
        let j = self.parent[i];
        if i == j {
            i
        } else {
            let k = self.find(j);
            self.parent[i] = k;
            k
        }
    }

    fn union(&mut self, mut i: usize, mut j: usize) {
        i = self.find(i);
        j = self.find(j);
        if i != j {
            self.parent[i] = j;
        }
    }
}

impl Solution {
    fn accounts_merge(accounts: Vec<Vec<String>>) -> Vec<Vec<String>> {
        let n = accounts.len();
        let mut btm: BTreeMap<&str, &str> = BTreeMap::new();
        for i in 0..n {
            let m = accounts[i].len();
            let name = &accounts[i][0];
            for j in 1..m {
                let email = &accounts[i][j];
                btm.insert(email, name);
            }
        }
        let count = btm.len();
        let mut uf = UnionFind::new(count);
        let mut owners = vec![];
        let mut emails = vec![];
        let mut ids = HashMap::new();
        for (i, (email, name)) in btm.into_iter().enumerate() {
            emails.push(email.to_string());
            owners.push(name.to_string());
            ids.insert(email.to_string(), i);
        }

        for i in 0..n {
            let m = accounts[i].len();
            for j in 2..m {
                let email_a = &accounts[i][j - 1];
                let email_b = &accounts[i][j];
                let id_a = ids[email_a];
                let id_b = ids[email_b];
                uf.union(id_a, id_b);
            }
        }
        let mut res: Vec<Vec<String>> = vec![];
        let mut hm: HashMap<usize, BTreeSet<usize>> = HashMap::new();
        for i in 0..count {
            let group_id = uf.find(i);
            hm.entry(group_id).or_default().insert(i);
        }
        for (group_id, ids) in hm.into_iter() {
            let mut v: Vec<String> = vec![];
            v.push(owners[group_id].clone());
            for id in ids {
                v.push(emails[id].to_string());
            }
            res.push(v);
        }
        res
    }
}

#[test]
fn test() {
    let accounts = vec![
        vec_string!["John", "johnsmith@mail.com", "john00@mail.com"],
        vec_string!["John", "johnnybravo@mail.com"],
        vec_string!["John", "johnsmith@mail.com", "john_newyork@mail.com"],
        vec_string!["Mary", "mary@mail.com"],
    ];
    let mut ans = vec![
        vec_string![
            "John",
            "john00@mail.com",
            "john_newyork@mail.com",
            "johnsmith@mail.com"
        ],
        vec_string!["John", "johnnybravo@mail.com"],
        vec_string!["Mary", "mary@mail.com"],
    ];
    let mut res = Solution::accounts_merge(accounts);
    res.sort_unstable();
    ans.sort_unstable();
    assert_eq!(ans, res);
}