What time and space complexity am I using with memoized countUnivalTrees problem

user_9885901 0
Viewed 47

I was attempting this challenge but took too long to execute for some inputs. The question is as follows

You are given a binary tree. Return the count of unival sub-trees in the given binary tree. In unival trees, all the nodes, below the root node, have the same value as the data of the root.

My first attempt failed because of the time complexity

def countUnivalTrees(tree):
    if tree is None:
        return 0

    leftCount = countUnivalTrees(tree.left)
    rightCount = countUnivalTrees(tree.right)

    count = leftCount + rightCount

    if isUnivalTree(tree, tree.value):
        count += 1

    return count


def isUnivalTree(tree, value):
    if tree is None:
        return True

    if tree.value != value:
        return False

    if tree.left is None and tree.right is None:
        return True

    return isUnivalTree(tree.left, value) and isUnivalTree(tree.right, value)

It gives the right answer but can take too long to execute for some given inputs.

My Second approach was to memoize what I've seen already.

def countUnivalTrees(root, memoized={}):
    if root is None:
        return 0

    memoized[root] = False

    leftCount = countUnivalTrees(root.left, memoized)
    rightCount = countUnivalTrees(root.right, memoized)

    count = leftCount + rightCount

    if root.left is not None and root.right is not None:
        if memoized[root.left] and memoized[root.right] and root.left.value == root.right.value and root.value == root.left.value:
            memoized[root] = True
            count += 1
    elif root.left is not None:
        if memoized[root.left] and root.value == root.left.value:
            memoized[root] = True
            count += 1
    elif root.right is not None:
        if memoized[root.right] and root.value == root.right.value:
            memoized[root] = True
            count += 1
    else:
        if isUnivalTree(root, root.value):
            memoized[root] = True
            count += 1

    return count

The isUnivalTree is indentical to the first attempt so I omitted it to not add the extra code. I was able to pass all tests. My plan was to have a time complexity of O(n) and O(n) space. I am having trouble understanding exactly what my complexity would be and if my approach is correct. By correct I mean working exactly how I believe it should be:

step 1: visit each node in the tree

step 2: from leaf to root check current root is unival or not. If it is save the root in memoized and map it to true.

step 3: if current root node has a left and right child node and that node in memoized maps to true meaning it is a unival tree than just check right and left child values to be the same and root value to be the same and is a unival tree memoize the current root and map to true. (also take into account for roots that have only one left child or only one right child and perform those checks as well O(1) operations)

So I'm imagining a big tree with many spanning left and right but with all the same values. Is it true that we are looking at all the nodes once O(n), where n is the number of nodes and then as we pop off the stack memoizing that current node and mapping to true if isUnivalTree and so on.

For some reason I feel like I might be missing a case where I should check if memoized[root.left] = False and if it is then not count it. Not sure really, I never memoized before besides for the fibonacci coding problem so this is new to me.

memoization dynamic-programming tree algorithm python
1 Answers

We can't see the memoization because we don't see your isUnivalTree. But if it is fast, then it is certainly O(n) time because you're only visiting each node a maximum of 2x (once to count trees, once to see if it was marked unival), on one of which you memoize it.

But there is no need to memoize. Just return all of the context that you need in one go. And if that is messier than you want to use, then have a trivial wrapper.

def unival_trees (root):
    return _unival_trees_and_meta(root)[0]

def _unival_trees_and_meta(root):
    if root.left is None:
        if root.right is None:
            return (1, True, root.value)
        else:
            count_right, right_is_unival, right_value = _unival_trees_and_meta(root.right)
            if right_is_unival and root.value == right_value:
                return (count_right + 1, True, root.value)
            else:
                return (count_right, False, root.value)
    else:
        count_left, left_is_unival, left_value = _unival_trees_and_meta(root.left)
        if root.right is None:
            if left_is_unival and left_value == root.value:
                return (count_left + 1, True, root.value)
            else:
                return (count_left, False, root.value)
        else:
            count_right, right_is_unival, right_value = _unival_trees_and_meta(root.right)
            if left_is_unival and right_is_unival and left_value == right_value and left_value == root.value:
                return (count_right + count_left + 1, True, root.value)
            else:
                return (count_right + count_left, False, root.value)
user_5854110
Related