Count the number of matching nodes in the subtree rooted at some node n.


from collections import namedtuple

Node = namedtuple('Node', 'key, left, right')

def build_tree(inorder, preorder):
    """Builds a Binary Tree from given inorder and preorder traversal.
    1. Take the first element from the preorder list.
    2. Get the index of the element from inorder list.
    3. Set the element as root.
    4. Take all elements left to root from inorder list.
    5. Take all elements right to root from inorder list.
    6. Calculate preorder list for left and right subtree respectively.
    7. Repeat from step 1 for each subtree.
    if not preorder:
    if len(preorder) == 1:
        return Node(preorder, None, None)
    root = preorder[0]
    root_index    = inorder.find(root)
    inorder_left  = inorder[:root_index]
    inorder_right = inorder[root_index+1:]
    preorder_left = preorder[1:len(inorder_left)+1]
    preorder_right = preorder[root_index+1:]
    return Node(
        build_tree(inorder_left, preorder_left),
        build_tree(inorder_right, preorder_right))    

def match_count(root, key1, key2):
    """Counts the number of node that matches key1 or key2 in the subtree.""" 
    if root is None:
        return 0
    count_left = match_count(root.left, key1, key2)
    count_right = match_count(root.right, key1, key2)
    if root.key == key1 or root.key == key2:
        return count_left + count_right + 1
    return count_left + count_right

root = build_tree('DBEACF', 'ABDECF')

import unittest

class TestMatchCount(unittest.TestCase):
    def setUp(self):
               / \
              B   C
             /\    \
            D  E    F
        root = build_tree('DBEACF', 'ABDECF')

    def test_noNodesFound(self):
        self.assertEqual(match_count(root, 'Z', 'Q'), 0)

    def test_nodesOnSameSide(self):
        self.assertEqual(match_count(root, 'D', 'E'), 2)

    def test_nodesOnDiffSide(self):
        self.assertEqual(match_count(root, 'D', 'F'), 2)

    def test_oneNodeIsRoot(self):
        self.assertEqual(match_count(root, 'D', 'A'), 2)

if __name__ == '__main__':

Time complexity

Since we check each node at least once hence it should be \$O(n)\$, where n is the number of nodes in the Binary Tree.


while unit testing I felt that I am not able to express intent clearly for my test case unlike Jasmine where we can write proper description in it(...) format.

  • \$\begingroup\$ For the testing you could search for "fluent" in the descriptions of the testing libraries, e.g. a quick search yielded assertpy and well-behaved-python. \$\endgroup\$ – ferada Oct 26 '15 at 17:26

I'll leave the comments on build_tree() as they stand in my other answer, and focus on your match_count() method:

  • Bad parameters for method? – It feels a little constructed, and special, to test for two key values. Why not one, or three, or an arbitrary number? Or a range?
  • Not optimal method name? – To me, it sounds a little strange with match_count. I just think of safety matches... I think I would prefer just count() (or count_matching()). Do note that when using count() we kind of hide/reimplement the default count() of Python. If your binary tree was within a class (as in my test cases) that could work nicely, but outside of classes it should probable be avoided and you should rather use the count_matching or match_count name.
  • Consider use of local variables – I don't oppose to using local variables like count_left and count, but in this rather simple case you could actually do:

    def count(node, key1, key2):
        """Counts the number of node that matches key1 or key2 in the subtree.""" 
        if node is None:
            return 0
        return ( (node.key == key1 or node.key == key2)
                 + match_count(node.left, key1, key2)
                 + match_count(node.right, key1, key2) )

Use predicates to generalize method

If you look at the accepted answer of "Most pythonic way of counting matching elements in something iterable", they suggest some neat ways of handling your issue at hand. In fact your match_count() can be replaced with the itertools.quantify recipe.

def quantify(seq, pred=None):
    "Count how many times the predicate is true in the sequence"
    return sum(imap(pred, seq))

You could either use this directly, and reuse any of the order traversals, like in:

def count(node, pred=None):
    """Count how many times the predicate is true in binary tree."""
    return sum(imap(pred, inorder(node)))

Another version would be to use the predicate similar to the first count() version above:

def count(node, pred=None):
    """Count how many times the predicate is true in binary tree."""
    if node is None:
        return 0

    return (not pred or pred(node.key)) + count(node.left, pred) + count(node.right, pred)

Have added a not pred in front of the test to allow for using this method with default parameter of None in which case it ends up counting all elements in binary tree.

This general approach opens up for counting whatever you feel like, but you need to call it slightly different. Here are some examples:

# Two variants to test against two different key values 
# Note that the lambda function will receive the node.key value
count(root, lambda x: x == 'D' or x=='E')
count(root, lambda x: x in 'DE')

# or use a designated method (accepting
def is_vowel(char):
    return char in 'AEIOU'
count(root, is_vowel)  

Complexity considerations

Time complexity of any of these implementations (including your original one) will be \$O(n)\$. Slightly different memory complexity, but neither algorithm will use much extra memory (with the exception when using a list implementation of inorder() which duplicates the binary tree into a list thusly requiring more memory).

| improve this answer | |
  • In general, too many local variables. It easily gets quite hard to follow how the variable has been computed, and how it is being used.

  • Building the tree.

    • Since preorder[root_index+1:] goes to preorder_right, it is reasonable to assume that preorder_left shall take the rest (except preorder[0] of course). That is,

      preorder_left = preorder[1:root_index + 1]

      The call to len is redundant.

    • As for local variables, I prefer

      return Node(root,
          build_tree(inorder[:root_index], preorder[1:root_index + 1]),
          build_tree(inorder[root_index+1:], preorder[root_index+1:]))
    • Special-casing len(preorder) == 1 is not necessary.

  • Search

    The returns are a bit confusing.

        count = match_count(root.left, key1, key2)
        count += match_count(root.right, key1, key2)
        if root.key == key1 or root.key == key2:
            count += 1
        return count

    looks cleaner. Notice the absence of redundant locals.

| improve this answer | |
  • \$\begingroup\$ The use of descriptive names increases readability according to me. \$\endgroup\$ – CodeYogi Oct 31 '15 at 4:59
  • \$\begingroup\$ @CodeYogi There is a fine line between being descriptive and being noise. \$\endgroup\$ – vnp Oct 31 '15 at 6:17

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.