diff --git a/problems/product.py b/problems/product.py new file mode 100644 index 0000000..48e3eb4 --- /dev/null +++ b/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!")