Trying to find nodes in binary tree

Question

I am trying to find nodeA and nodeB in a binary tree branch. I am returning a sum that should equal the number of nodes found. Here is my implementation

def checkSubtree(self, nodeA, nodeB, node):
        if node is None:
            return 0
        if node is nodeA or node is nodeB:
            return 1
        return self.checkSubtree(nodeA, nodeB, node.leftChild) + self.checkSubtree(nodeA, nodeB, node.rightChild)

I keep getting 1 when i should be getting 2. I figured out that its because the second if statement is executed on the first pass and returns immediately.

How do I improve this. I do not want to use an extra variable to store the result and return that.

Eric Duminil · Accepted Answer · 2017-07-13 12:18:57Z

1

nodeA or nodeB nodes could still have children, which could possibly be nodeA or nodeB nodes.

With return 1, you stopped your recursive search as soon as you encountered a nodeA or nodeB node.

One solution would be:

def checkSubtree(self, nodeA, nodeB, node):
    if node is None:
        return 0
    else:
        children_sum = self.checkSubtree(nodeA, nodeB, node.leftChild) + self.checkSubtree(nodeA, nodeB, node.rightChild)
        if node is nodeA or node is nodeB:
            return 1 + children_sum # <-- Important. Not just 1 !
        else:
            return children_sum

If for whatever reason you do now want to define a local variable:

def checkSubtree(self, nodeA, nodeB, node):
    if node is None:
        return 0
    else:
        return self.checkSubtree(nodeA, nodeB, node.leftChild) + self.checkSubtree(
            nodeA, nodeB, node.rightChild) + int(node is nodeA or node is nodeB)

edited Jul 13, 2017 at 12:18

answered Jul 13, 2017 at 12:07

Eric Duminil

54.6k10 gold badges80 silver badges134 bronze badges

Sign up to request clarification or add additional context in comments.

6 Comments

redixhumayun Over a year ago

Can you explain that with code? I didn't follow you

redixhumayun Over a year ago

The last part of my post said "I do not want to use an extra variable to store the result and return that". So I am assuming theres no way to do this without using an extra variable?

cs95 Over a year ago

@ZaidHumayun Why would you want to use an extra variable when you don't need to.

Eric Duminil Over a year ago

@ZaidHumayun: Having a correct answer seems more important to me than not using an extra variable. Still, I'll update the answer.

redixhumayun Over a year ago

@cᴏʟᴅsᴘᴇᴇᴅ - I said I don't want to use an extra variable.

|

Collectives™ on Stack Overflow

Trying to find nodes in binary tree

1 Answer 1

6 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

6 Comments

Your Answer

Sign up or log in

Post as a guest

Related