title: Count number of leaves in a tree - Failed edge case?
body: <p>In an online assessment I was asked to count the number of leaves in a tree. The tree is given in parent-array representation, meaning the tree has n nodes with labels 0, 1, 2, .., n-1, and you are passed a length n array p, where p[i] returns the label of the parent of node i, except when i is the root of the tree in which case p[i] is -1. </p>
  
  <p>I guess one thing to note is that the problem was as stated above, so there were no extra conditions such as e.g. it being a binary tree.</p>
  
  <p>I thought this was a fairly straight forward problem, but the code that I submitted failed a "Small Tree Case" on the testing platform (which does not let you see the test cases). It passed the other tests, including a performance test on a large tree. I've thought about it for a while but I still can not see what the flaw in my algorithm or handling of some edge case is. I guess one thing to note is that the problem was as stated above, so there were no extra conditions such as e.g. it being a binary tree.</p>
  
  <pre class="lang-py prettyprint-override"><code>def countLeaves(p):
      n = len(p)
      if p is None or n == 0 : return 0
      if n == 1 or n == 2 : return 1
  
      leaves = set(range(n))
      for i in range(n):
          if p[i] == -1: # i is root of tree with &gt;1 node, can't be a leaf
              leaves.discard(i)
          else: # p[i] is parent of node i, can't be a leaf
              leaves.discard(p[i])
  
      return len(leaves)
  </code></pre>
  
  <p>In trying to fix the failed "Small tree case" I also tried returning None if p is None, returning None if n == 0, or both modifications together, but to no success. If anyone could point out what the error in my code may have been I would greatly appreciate it. Thank you.</p>