Binary Trees

 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
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
 
import sys
import unittest
from dataclasses import dataclass
from typing import *
sys.setrecursionlimit(10**6)

## binary trees are another example of mixed compound data:

BinTree = Union['BTNode',None]

@dataclass(frozen=True)
class BTNode:
    val : Any
    left : BinTree
    right : BinTree

# examples of data
bt1 = None
bt2 = BTNode(4,None,None)
bt3 = BTNode(4,
             BTNode("thingy",None,None),
             None)
bt4 = BTNode("swim",
             BTNode("run",BTNode("walk",None,None), BTNode("banana",None,None)),
             BTNode("sleep",None,None))




# count the nodes in a binary tree:
def bt_nodes(bt : BinTree) -> int:
    match bt:
        case None:
            return 0
        case BTNode(_,l,r):
            return bt_nodes(l) + bt_nodes(r) + 1



# does the list contain the string "banana" ?
def contains_banana(bt : BinTree) -> bool:
    match bt:
        case None:
            return False
        case BTNode(v,l,r):
            return (v == "banana") or contains_banana(l) or contains_banana(r)

# flip the binary tree left to right
def mirror(bt : BinTree) -> BinTree:
    match bt:
        case None:
            return None
        case BTNode(v,l,r):
            return BTNode(v,mirror(r),mirror(l))

class MyTests(unittest.TestCase):

    def test_count_nodes(self):
        self.assertEqual(0,bt_nodes(bt1))
        self.assertEqual(1,bt_nodes(bt2))
        self.assertEqual(2,bt_nodes(bt3))
        self.assertEqual(5,bt_nodes(bt4))

    def test_contains_banana(self):
        self.assertEqual(False,contains_banana(bt1))
        self.assertEqual(True,contains_banana(bt4))
        self.assertEqual(False,contains_banana(bt2))

    def test_mirror(self):
        self.assertEqual(bt1,mirror(bt1))
        self.assertEqual(bt2,mirror(bt2))
        self.assertEqual(BTNode(4,
                                None,
                                BTNode("thingy",None,None)),
                         mirror(bt3))
        self.assertEqual(BTNode("swim",
                                BTNode("sleep",None,None),
                                BTNode("run",BTNode("banana",None,None), BTNode("walk",None,None))),
                         mirror(bt4))