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


  • 49

    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.


  • 1

    @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...


  • 0

    @StefanPochmann Hi, I found the second DFS method doesn't pass the test. I printed out the neighbors before and after visit(0) and found that visit doesn't do the recursive.


Log in to reply
 

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