FilesExpand file tree

 master

/

depth_sum.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

66 lines (53 loc) · 1.58 KB

 master

/

depth_sum.py

Copy path

Top

File metadata and controls

• Code
• Blame

66 lines (53 loc) · 1.58 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

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

"""

Write a function depthSum returns the sum of the values stored

in a binary search tree of integers weighted by the depth of each value.

For example:

9

/ \

6 12

/ \ / \

3 8 10 15

/ \

7 18

depth_sum = 1*9 + 2*(6+12) + 3*(3+8+10+15) + 4*(7+18)

"""

import unittest

from bst import Node

from bst import bst

def depth_sum(root, n):

if root:

return recur_depth_sum(root, 1)

def recur_depth_sum(root, n):

if root is None:

return 0

elif root.left is None and root.right is None:

return root.data * n

else:

return n * root.data + recur_depth_sum(root.left, n+1) + recur_depth_sum(root.right, n+1)

"""

The tree is created for testing:

9

/ \

6 12

/ \ / \

3 8 10 15

/ \

7 18

depth_sum = 1*9 + 2*(6+12) + 3*(3+8+10+15) + 4*(7+18)

"""

class TestSuite(unittest.TestCase):

def setUp(self):

self.tree = bst()

self.tree.insert(9)

self.tree.insert(6)

self.tree.insert(12)

self.tree.insert(3)

self.tree.insert(8)

self.tree.insert(10)

self.tree.insert(15)

self.tree.insert(7)

self.tree.insert(18)

def test_depth_sum(self):

self.assertEqual(253, depth_sum(self.tree.root, 4))

if __name__ == '__main__':

unittest.main()