Coverage for blog/dsa/leetcode/bt/is_balanced/__init__.py: 72%

1from typing import Any 

2 

3import pytest 

4 

5from dsa.leetcode.bt.from_desc import TreeNode 

6 

7# NOTE: https://leetcode.com/problems/balanced-binary-tree/description/ 

8# NOTE: https://leetcode.com/problems/binary-tree-inorder-traversal/description/ 

9# NOTE: https://leetcode.com/problems/balance-a-binary-search-tree/editorial/ 

10 

11 

12# start snippit solution_initial 

13class SolutionInitial: 

14 def _isBalanced( 

15 self, root: TreeNode | None, *, _depth: int = 0 

16 ) -> tuple[bool, int]: 

17 if root is None: 

18 return True, _depth 

19 

20 left_is_balanced, left_depth = self._isBalanced(root.left, _depth=_depth + 1) 

21 right_is_balanced, right_depth = self._isBalanced(root.right, _depth=_depth + 1) 

22 

23 self_is_balanced = ( 

24 abs(left_depth - right_depth) <= 1 

25 and right_is_balanced 

26 and left_is_balanced 

27 ) 

28 depth = max(left_depth, right_depth) 

29 

30 return self_is_balanced, depth 

31 

32 def isBalanced(self, root: TreeNode | None) -> bool: 

33 is_balanced, _ = self._isBalanced(root) 

34 return is_balanced 

35 

36 # end snippit solution_final 

37 

38 

39# start snippit solution 

40class Solution: 

41 def _isBalanced(self, root: TreeNode | None, *, _depth: int = 0) -> int: 

42 if root is None: 

43 return _depth 

44 

45 _depth += 1 

46 left_depth = self._isBalanced(root.left, _depth=_depth) 

47 if left_depth < 0: 

48 return -1 

49 

50 right_depth = self._isBalanced(root.right, _depth=_depth) 

51 if right_depth < 0: 

52 return -1 

53 

54 if abs(left_depth - right_depth) > 1: 

55 return -1 

56 

57 return max(left_depth, right_depth) 

58 

59 def isBalanced(self, root: TreeNode | None) -> bool: 

60 return self._isBalanced(root) >= 0 

61 

62 # end snippit solution 

63 

64 

65@pytest.fixture 

66def solution(): 

67 return Solution() 

68 

69 

70@pytest.mark.parametrize( 

71 "question, is_balanced", 

72 ( 

73 ({"value": 1}, True), 

74 ( 

75 { 

76 "value": 1, 

77 "right": { 

78 "value": 2, 

79 "right": { 

80 "value": 3, 

81 }, 

82 }, 

83 }, 

84 False, 

85 ), 

86 ( 

87 { 

88 "value": 2, 

89 "right": {"value": 3}, 

90 "left": {"value": 1}, 

91 }, 

92 True, 

93 ), 

94 ( 

95 { 

96 "value": 1, 

97 "left": {"value": 2}, 

98 "right": { 

99 "value": 4, 

100 "left": { 

101 "value": 5, 

102 "right": {"value": 6}, 

103 "left": {"value": 7}, 

104 }, 

105 "right": {"value": 8}, 

106 }, 

107 }, 

108 False, 

109 ), 

110 ), 

111) 

112def test_solution_is_balanced( 

113 solution: Solution, question: dict[str, Any], is_balanced: bool 

114): 

115 

116 question_tree = TreeNode.fromDict(question) 

117 is_balanced_computed = solution.isBalanced(question_tree) 

118 print(is_balanced_computed, is_balanced) 

119 assert is_balanced_computed == is_balanced