LeetCode - 1. Two Sum, python3, Accepted

solution.py

@@ -11,12 +11,8 @@ def twoSum(nums, target):
 		:type target: int
 		:rtype: List[int]
 		"""
-		a = 0
-		b = 0
+		cache = {}
 		for idx, num in enumerate(nums):
-			for idx2, num2 in enumerate(nums[idx + 1:]):
-				if num2 == target - num:
-					a = idx
-					b = idx2 + idx + 1
-					break
-		return [a, b]
+			if target - num in cache.keys():
+				return [cache[target - num], idx]
+			cache[num] = idx

solution_test.py

@@ -11,9 +11,10 @@ def test_sum(self):
         self.assertEqual(1, Solution.sub(3, 2))
         self.assertEqual(-1, Solution.sub(2, 3))
 
-    def test_sum(self):
+    def test_twoSum(self):
         self.assertEqual([0, 1], Solution.twoSum([2, 7, 11, 15], 9))
         self.assertEqual([1, 2], Solution.twoSum([3, 2, 4], 6))
+        self.assertEqual([0, 1], Solution.twoSum([3, 3], 6))
 
 
 if __name__ == '__main__':