Update 366. Find Leaves of Binary Tree.md

366. Find Leaves of Binary Tree.md

@@ -18,7 +18,7 @@
 
 AC代码
 
-```
+```python
 class Solution(object):
     def findLeaves(self, root):
         """
@@ -26,29 +26,27 @@ class Solution(object):
         :rtype: List[List[int]]
         """
         def findLeaf(root):
-        	if root == None:
-        		return []
-        	elif root.left == None and root.right == None:
-        		return [root.val]
-        	else:
-        		return findLeaf(root.left)  + findLeaf(root.right)
-
+            if not root:
+                return []
+            elif not root.left and not root.right:
+                return [root.val]
+            else:
+                return findLeaf(root.left) + findLeaf(root.right)
         def removeLeaf(root):
-        	if root == None:
-        		return None
-        	elif root.left == None and root.right == None:
-        		return None
-        	else:
-        		if root.left:
-        			root.left = removeLeaf(root.left)
-        		if root.right:
-        			root.right = removeLeaf(root.right)
-        		return root 
-        
+            if not root:
+                return None
+            elif not root.left and not root.right:
+                return None
+            else:
+                if root.left:
+                    root.left = removeLeaf(root.left)
+                if root.right:
+                    root.right = removeLeaf(root.right)
+                return root
         res = []
         while root:
-        	res.append(findLeaf(root))
-        	root = removeLeaf(root)
+            res.append(findLeaf(root))
+            root = removeLeaf(root)
         return res
 ```