Python solution Based on BFS and DFS


  • 1
    G

    '''

    def isBalanced(self, root):
        if not root: return True
        levels=[[root]]
        parrent={}
        q=collections.deque()
        q.append((root,0))
        while q:
            node, k = q.popleft()
            if node.left:
                q.append((node.left, k+1))
                parrent[node.left]=node
                try:
                    levels[k+1].append(node.left)
                except:
                    levels.append([node.left])
            if node.right:
                q.append((node.right, k+1))
                parrent[node.right]=node
                try:
                    levels[k+1].append(node.right)
                except:
                    levels.append([node.right])
        levels.reverse()        # traverse nodes from bottom
        
        height={u:0 for level in levels for u in level}
        for i in range(len(levels)-1):
            for u in levels[i]:
                p=parrent[u]
                height[p]=max(height[p], height[u]+1)
        
        q.append(root)
        while q:
            node=q.popleft()
            if height[node]<=1: continue
            if not node.left or not node.right: return False
            if abs(height[node.left]-height[node.right])>1: return False
            q.append(node.left)
            q.append(node.right)
        return True
    

    '''
    Firstly, I do a level order traversal of the tree. Then I calculate the height of each node, the height of leaves is treated as 0. Finally I do a tree traversal again and use the height value to check if the tree is balanced.

    The running time for this BFS algorithm is around 92 ms. It is relatively long because it goes through all the nodes three times.

    A DFS method should be faster as it only needs to traverse the tree only once:
    '''

    def isBalanced(self, root):
        if not root: return True
        qstack=[root]      # Used queue stack for DFS of the tree
        levels={None: 0}   # Store the levels/height of each subtree, a None subtree is treated as 0-height tree.
        while qstack:
            node=qstack[-1]
            if (not node.left) and (not node.right):
                levels[node]=1
                qstack.pop()
                continue
            if (node.left in levels) and (node.right in levels): # Note that an empty subtree is always in levels.
                if abs(levels[node.left]-levels[node.right])>1:
                    return False
                else:
                    levels[node]=1+max(levels[node.left], levels[node.right])
                    qstack.pop()
            else:
                if node.right:
                    qstack.append(node.right)
                if node.left:
                    qstack.append(node.left)
        return True
    

    '''

    The running time is around 80 ms.


Log in to reply
 

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