-
Notifications
You must be signed in to change notification settings - Fork 39
/
solution_08.py
30 lines (22 loc) · 883 Bytes
/
solution_08.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
def coding_problem_08(bt):
"""
A unival tree (which stands for "universal value") is a tree where all nodes have the same value.
Given the root to a binary tree, count the number of unival subtrees.
Example:
>>> btree = (0, (0, (0, None, None), (0, (0, None, None), (0, None, None))), (1, None, None))
>>> coding_problem_08(btree)
6
"""
def unival_count(btree):
val, ln, rn = btree
if ln is None and rn is None: # leaf case
return 1, True, val
lcount, is_uni_l, lval = unival_count(ln)
rcount, is_uni_r, rval = unival_count(rn)
is_unival = is_uni_l and is_uni_r and val == lval and val == rval
count = lcount + rcount + is_unival
return count, is_unival, val
return unival_count(bt)[0]
if __name__ == '__main__':
import doctest
doctest.testmod(verbose=True)