题目

https://leetcode.com/problems/accounts-merge/

题解

HashMap 版的并查集。参考了：leetcode 684. Redundant Connection | 684. 冗余连接（并查集）

class Solution {
    public List<List<String>> accountsMerge(List<List<String>> accounts) {
        Map<String, String> parent = new HashMap<>();
        Map<String, String> name = new HashMap<>();
        Map<String, Integer> size = new HashMap<>();
        // init
        for (List<String> list : accounts) {
            name.put(list.get(1), list.get(0));
            parent.put(list.get(1), list.get(1));
            size.put(list.get(1), 1);
            for (int i = 2; i < list.size(); i++) {
                name.put(list.get(i), list.get(0));
                parent.put(list.get(i), list.get(i));
                size.put(list.get(i), 1);
            }
        }
        // union find
        for (List<String> list : accounts) {
            for (int i = 2; i < list.size(); i++) {
                String r0 = findRoot(parent, list.get(1)); // r0会被修改，不能放在外层不变！！
                String r1 = findRoot(parent, list.get(i));
                if (r0.equals(r1)) continue;
                if (size.get(r0) > size.get(r1)) {
                    parent.put(r1, r0);
                } else {
                    parent.put(r0, r1);
                }
                size.put(r0, size.get(r0) + size.get(r1));
                size.put(r1, size.get(r0) + size.get(r1));
            }
        }
        HashMap<String, Set<String>> resMap = new HashMap<>();
        for (List<String> list : accounts) {
            for (int i = 1; i < list.size(); i++) {
                String root = findRoot(parent, list.get(i));
                resMap.computeIfAbsent(root, k -> new TreeSet<>());
                resMap.get(root).add(list.get(i));
            }
        }
        List<List<String>> result = new ArrayList<>();
        for (String root : resMap.keySet()) {
            List<String> list = new ArrayList<>();
            list.add(name.get(root));
            list.addAll(resMap.get(root));
            result.add(list);
        }
        return result;
    }

    // 从s开始一直往上找，直到不能再往上，找到代表节点，返回
    public String findRoot(Map<String, String> parent, String s) {
        String cur = s;
        while (!cur.equals(parent.get(cur))) {
            cur = parent.get(cur);
        }
        while (!s.equals(parent.get(s))) { // 路径压缩
            String next = parent.get(s);
            parent.put(s, cur);
            s = next;
        }
        return cur;
    }
}