8-10 lines, Union-Find, DFS and BFS


  • 44

    There are so many different approaches and so many different ways to implement each. I find it hard to decide, so here are several :-)

    In all of them, I check one of these tree characterizations:

    • Has n-1 edges and is acyclic.
    • Has n-1 edges and is connected.

    Solution 1 ... Union-Find

    The test cases are small and harmless, simple union-find suffices (runs in about 50~60 ms).

    def validTree(self, n, edges):
        parent = range(n)
        def find(x):
            return x if parent[x] == x else find(parent[x])
        def union(xy):
            x, y = map(find, xy)
            parent[x] = y
            return x != y
        return len(edges) == n-1 and all(map(union, edges))
    

    A version without using all(...), to be closer to other programming languages:

    def validTree(self, n, edges):
        parent = range(n)
        def find(x):
            return x if parent[x] == x else find(parent[x])
        for e in edges:
            x, y = map(find, e)
            if x == y:
                return False
            parent[x] = y
        return len(edges) == n - 1
    

    A version checking len(edges) != n - 1 first, as parent = range(n) could fail for huge n:

    def validTree(self, n, edges):
        if len(edges) != n - 1:
            return False
        parent = range(n)
        def find(x):
            return x if parent[x] == x else find(parent[x])
        def union(xy):
            x, y = map(find, xy)
            parent[x] = y
            return x != y
        return all(map(union, edges))
    

    Solution 2 ... DFS

    def validTree(self, n, edges):
        neighbors = {i: [] for i in range(n)}
        for v, w in edges:
            neighbors[v] += w,
            neighbors[w] += v,
        def visit(v):
            map(visit, neighbors.pop(v, []))
        visit(0)
        return len(edges) == n-1 and not neighbors
    

    Or check the number of edges first, to be faster and to survive unreasonably huge n:

    def validTree(self, n, edges):
        if len(edges) != n - 1:
            return False
        neighbors = {i: [] for i in range(n)}
        for v, w in edges:
            neighbors[v] += w,
            neighbors[w] += v,
        def visit(v):
            map(visit, neighbors.pop(v, []))
        visit(0)
        return not neighbors
    

    For an iterative version, just replace the three "visit" lines with

        stack = [0]
        while stack:
            stack += neighbors.pop(stack.pop(), [])
    

    Solution 3 ... BFS

    Just like DFS above, but replace the three "visit" lines with

        queue = [0]
        for v in queue:
            queue += neighbors.pop(v, [])
    

    or, since that is not guaranteed to work, the safer

        queue = collections.deque([0])
        while queue:
            queue.extend(neighbors.pop(queue.popleft(), []))

  • 0
    T

    Very elegant. Thanks for providing them.


  • 0
    C

    Can anyone rewrite this in C++? Thanks


  • 1
    H

    Hello, thanks for your code. Here is the C++ version of union-find algorithm.

    class Solution {
    public:
        bool validTree(int n, vector<pair<int, int>>& edges) {
            if(edges.size() != n-1) return false;
    
            parent.resize(n);
            iota(parent.begin(), parent.end(), 0);
            
            for(auto e : edges)
            {
                auto x = find(e.first), y = find(e.second);
                if(x == y) return false;
                parent[x] = y;
            }
            
            return true;
        }
    private:
        vector<int> parent;
        int find(int x)
        {
            return parent[x] == x ? x : find(parent[x]);
        }
    };

  • 0
    L

    what's the time complexity of Union Find solution?


  • 0

    Stefan, why don't you do the following in the iterative DFS solution:

    stack = [0]
    while stack and neighbors:
        stack += neighbors.pop(stack.pop(), [])
    

    If neighbors becomes empty, what's the point of continuing the DFS?


  • 0

    @agave I don't remember. I suspect I just didn't see it. Though it's also possible that I did see it and didn't like it because it's more code and doesn't improve the complexity class.


  • 0

    @StefanPochmann I agree it doesn't improve the complexity class, but it is still an optimization that comes with just 13 more characters of code :-)


  • 0
    F

    @hsi-hung When you doing the find operation, I think it is more efficient to do the path compression.


  • 0
    F

    @liuxiyun.nku I think by using path compression, the time complexity can be O(n). Without path compression, the time complexity is O(n^2)


  • 0
    C

    @StefanPochmann said in 8-10 lines, Union-Find, DFS and BFS:

        queue = [0]
        for v in queue:
            queue += neighbors.pop(v, [])
    

    or, since that is not guaranteed to work, the safer

    Why do you say that is not guaranteed to work?


  • 0

    @ccyjoshua Well, in general things can go wrong if you modify something while you're iterating it, for example this: https://stackoverflow.com/q/6260089/1672429
    I do think that extending a list during iteration is or at least should be safe, but I'm not aware of the documentation stating this. If you know/find something about it, I'd be interested.


  • 0
    C

    @StefanPochmann Thanks! I see what you are saying. That's why I got ConcurrentModificationException when I was trying to rewrite this in Java. I am not familiar with the extend in Python...


Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.