Completed product question

problems/product.py

@@ -0,0 +1,43 @@
+"""
+Given an array of integers, return a new array such that each element at index i of the new array is the product of all the numbers in the original array except the one at i.
+
+For example, if the input is [1, 2, 3, 4, 5], the expected output is [120, 60, 40, 30, 24].
+
+Follow-up: What if you can't use division?
+"""
+
+def product(arr: list[int]) -> list[int]:
+    # first we calculate prefix/suffix. We add a spot with multiplicative identity (1) for convenience.
+    prefix = [1]*len(arr)
+    suffix = [1]*len(arr)
+
+    p = 1
+    for i in range(0, len(arr)-1):
+        p *= arr[i]
+        prefix[i+1] = p
+
+    p = 1
+    for i in range(len(arr)-1, 0, -1):
+        p *= arr[i]
+        suffix[i-1] = p
+
+    ans = [1]*len(arr)
+    for i in range(len(ans)):
+        ans[i] = prefix[i] * suffix[i]
+    return ans
+
+
+# --- tests ---
+
+assert product([2, 3, 4, 5]) == [60, 40, 30, 24]
+assert product([1, 2, 3, 4, 5]) == [120, 60, 40, 30, 24]
+assert product([1, 2]) == [2, 1]
+
+# zero cases
+assert product([0, 1, 2, 3]) == [6, 0, 0, 0]
+assert product([0, 0, 2, 3]) == [0, 0, 0, 0]
+
+# single element
+assert product([5]) == [1]
+
+print("All tests passed!")