CODEWITHSUNDEEP
python
DMan
DMan

Reputation: 61

Python: Count nodes of tree

Actually, I'm writing a method which returns the size of a tree:

def size(self):
    count = 0
    hasLeft, hasRight = self.left is not None, self.right is not None
    if hasLeft:
        count += self.left.size()
    if hasRight:
        count += self.right.size()
    if (hasLeft or hasRight) and self.root_value is not None:
        count += 1
    return count

The method works, but only for interior nodes :( I'm sure there must be a very easy solution to get the size of the tree...but how? Example: If I call:

tree=Tree_Class(2,Tree_Class(1,Tree_Class(3),Tree_Class(20)),Tree_Class(13,Tree_Class(33),Tree_Class(39)))
tree.size()

The desired output is: 7

Thanks for any help!

Upvotes: 1

Views: 3129

Answers (1)

Ying Xiong
Ying Xiong

Reputation: 4928

The problem is your statement if (hasLeft or hasRight) and self.root_value is not None: --- this check is wrong, it only +=1 if the node is not a leaf. When size is called, this node must be valid (according to your logic), and therefore count should always += 1.

A simpler code would look like the following

class Tree_Class:
    def __init__(self, t, left=None, right=None):
        self.root_value = t
        self.left = left
        self.right = right

    def size(self):
        count = 1
        if self.left:
            count += self.left.size()
        if self.right:
            count += self.right.size()
        return count

And the results:

>>> tree=Tree_Class(2,Tree_Class(1,Tree_Class(3),Tree_Class(20)),Tree_Class(13,Tree_Class(33),Tree_Class(39)))
>>> tree.size()
7

Upvotes: 3

Related Questions