fork download
  1. import math
  2. import sys
  3. import time
  4.  
  5. sys.setrecursionlimit(500000)
  6.  
  7. class TreeNode:
  8. def __init__(self, left=None, right=None):
  9. self.left = left
  10. self.right = right
  11.  
  12. class Solution:
  13. def isBalanced(self, root):
  14. if not root:
  15. return True
  16.  
  17. left = self.height(root.left)
  18. right = self.height(root.right)
  19.  
  20. if abs(left - right) > 1:
  21. return False
  22.  
  23. return self.isBalanced(root.left) and self.isBalanced(root.right)
  24.  
  25. def height(self, node):
  26. if not node:
  27. return 0
  28. return 1 + max(self.height(node.left), self.height(node.right))
  29.  
  30. def build_linked_list(n):
  31. head = None
  32. for _ in range(n):
  33. head = TreeNode(left=head)
  34. return head
  35.  
  36. def build_perfect_tree(n):
  37. if n <= 0:
  38. return None
  39. left_n = (n - 1) // 2
  40. right_n = n - 1 - left_n
  41. return TreeNode(left=build_perfect_tree(left_n), right=build_perfect_tree(right_n))
  42.  
  43. def build_worst_case_tree(k):
  44. """
  45. k is the length of the spine.
  46. """
  47. if k <= 0:
  48. return None
  49. root = TreeNode(left=build_worst_case_tree(k - 1))
  50. if k > 1:
  51. root.right = build_linked_list(k - 1)
  52. return root
  53.  
  54.  
  55. N = 200_000
  56. # Calculate the spine length k needed to get roughly N total nodes
  57. # k*(k+1)/2 = N => k^2 + k - 2N = 0 => k = (-1 +/- sqrt(1+8N))/2
  58. # We take the positive root and round down to get an integer k.
  59. # That's k = int((-1 + math.sqrt(1 + 8 * N)) / 2)
  60. k = int((-1 + math.sqrt(1 + 8 * N)) / 2)
  61.  
  62. trees = {
  63. "Linked List": build_linked_list(N),
  64. "Perfect Tree": build_perfect_tree(N),
  65. "Worst-Case Tree": build_worst_case_tree(k)
  66. }
  67.  
  68. def count_nodes(node):
  69. if not node:
  70. return 0
  71. return 1 + count_nodes(node.left) + count_nodes(node.right)
  72.  
  73. print("Node counts:")
  74. for name, tree in trees.items():
  75. print(f"{name: <16} : {count_nodes(tree)} nodes")
  76.  
  77. solver = Solution()
  78. print(f"\nTime it takes:")
  79.  
  80. for name, tree in trees.items():
  81. start = time.perf_counter()
  82. solver.isBalanced(tree)
  83. end = time.perf_counter()
  84.  
  85. print(f"{name: <16} : {end - start:.5f} seconds")
Time limit exceeded #stdin #stdout 5s 88592KB
stdin
Standard input is empty
stdout
Standard output is empty