Update sort-integers-by-the-power-value.py

Author: kamyu104

Python/sort-integers-by-the-power-value.py

@@ -38,19 +38,14 @@ def partition_around_pivot(left, right, pivot_idx, nums, compare):
                     left = new_pivot_idx + 1
 
         def power_value(x):
-            y = x
-            if x not in Solution.dp:
-                result = 0
-                while x > 1:
-                    result += 1
-                    if x%2:
-                        x = 3*x + 1
-                    else:
-                        x //= 2
-                    if x in Solution.dp:
-                        result += Solution.dp[x]
-                        break
-                Solution.dp[y] = result
+            y, result = x, 0
+            while x > 1 and x not in Solution.dp:
+                result += 1
+                if x%2:
+                    x = 3*x + 1
+                else:
+                    x //= 2
+            Solution.dp[y] = result + (Solution.dp[x] if x > 1 else 0)
             return Solution.dp[y], y
 
         arr = map(power_value, range(lo, hi+1))
@@ -71,19 +66,14 @@ def getKth(self, lo, hi, k):
         :rtype: int
         """
         def power_value(x):
-            y = x
-            if x not in Solution2.dp:
-                result = 0
-                while x > 1:
-                    result += 1
-                    if x%2:
-                        x = 3*x + 1
-                    else:
-                        x //= 2
-                    if x in Solution2.dp:
-                        result += Solution2.dp[x]
-                        break
-                Solution2.dp[y] = result
+            y, result = x, 0
+            while x > 1 and x not in Solution2.dp:
+                result += 1
+                if x%2:
+                    x = 3*x + 1
+                else:
+                    x //= 2
+            Solution2.dp[y] = result + (Solution2.dp[x] if x > 1 else 0)
             return Solution2.dp[y], y
 
         return sorted(range(lo, hi+1), key=power_value)[k-1]