Find the weak connected component in the directed graph

solution

http://www.lintcode.com/en/problem/find-the-weak-connected-component-in-the-directed-graph/

code

/**
 * Definition for Directed graph.
 * struct DirectedGraphNode {
 *     int label;
 *     vector<DirectedGraphNode *> neighbors;
 *     DirectedGraphNode(int x) : label(x) {};
 * };
 */

class UnionFind {
public:
    map<int, int> hash;

    void add(int key, int val)
    {
        hash[key] = val;
    }

    int find(int key)
    {
        while(key != hash[key])
        {
            key = hash[key];
        }
        return key;
    }

    int compress_find(int key)
    {
        if(key == hash[key])
        {
            return key;
        }

        int boss = find(key);


        while(key != boss)
        {
            int next = hash[key];
            hash[key] = boss;
            key = next;
        }

        return boss;
    }

    int union_father(int s1, int s2)
    {
        int father_s1 = find(s1);
        int father_s2 = find(s2);
        if(father_s1 != father_s2)
        {
            hash[father_s1] = father_s2;
        }
    }
};

class Solution {
public:
    /**
     * @param nodes a array of directed graph node
     * @return a connected set of a directed graph
     */
    vector<vector<int>> connectedSet2(vector<DirectedGraphNode*>& nodes) {
        // Write your code here
        vector<vector<int>> ret;
        if(nodes.size() == 0)
        {
            return ret;
        }

        UnionFind uf;

        for(DirectedGraphNode* tmp : nodes)
        {
            uf.add(tmp->label, tmp->label);
        }

        for(DirectedGraphNode* tmp : nodes)
        {
            for(DirectedGraphNode* p : tmp->neighbors)
            {
                uf.union_father(tmp->label, p->label);
            }
        }

        map<int, vector<int>> hash;
        for(DirectedGraphNode* tmp : nodes)
        {
            int boss = uf.compress_find(tmp->label);
            if(hash.find(boss) == hash.end())
            {
                vector<int> v;
                v.push_back(tmp->label);
                hash[boss] = v;
            }
            else
            {
                hash[boss].push_back(tmp->label);
            }
        }

        for(map<int, vector<int>>::iterator it = hash.begin(); it != hash.end(); it++)
        {
            ret.push_back(it->second);
        }

        return ret;

    }
};