Create product_sum.py

algorithm/data_structures/arrays/product_sum.py

@@ -0,0 +1,98 @@
+"""
+Calculate the Product Sum from a Special Array.
+reference: https://dev.to/sfrasica/algorithms-product-sum-from-an-array-dc6
+
+Python doctests can be run with the following command:
+python -m doctest -v product_sum.py
+
+Calculate the product sum of a "special" array which can contain integers or nested
+arrays. The product sum is obtained by adding all elements and multiplying by their
+respective depths.
+
+For example, in the array [x, y], the product sum is (x + y). In the array [x, [y, z]],
+the product sum is x + 2 * (y + z). In the array [x, [y, [z]]],
+the product sum is x + 2 * (y + 3z).
+
+Example Input:
+[5, 2, [-7, 1], 3, [6, [-13, 8], 4]]
+Output: 12
+
+"""
+
+
+def product_sum(arr: list[int | list], depth: int) -> int:
+    """
+    Recursively calculates the product sum of an array.
+
+    The product sum of an array is defined as the sum of its elements multiplied by
+    their respective depths.  If an element is a list, its product sum is calculated
+    recursively by multiplying the sum of its elements with its depth plus one.
+
+    Args:
+        arr: The array of integers and nested lists.
+        depth: The current depth level.
+
+    Returns:
+        int: The product sum of the array.
+
+    Examples:
+        >>> product_sum([1, 2, 3], 1)
+        6
+        >>> product_sum([-1, 2, [-3, 4]], 2)
+        8
+        >>> product_sum([1, 2, 3], -1)
+        -6
+        >>> product_sum([1, 2, 3], 0)
+        0
+        >>> product_sum([1, 2, 3], 7)
+        42
+        >>> product_sum((1, 2, 3), 7)
+        42
+        >>> product_sum({1, 2, 3}, 7)
+        42
+        >>> product_sum([1, -1], 1)
+        0
+        >>> product_sum([1, -2], 1)
+        -1
+        >>> product_sum([-3.5, [1, [0.5]]], 1)
+        1.5
+
+    """
+    total_sum = 0
+    for ele in arr:
+        total_sum += product_sum(ele, depth + 1) if isinstance(ele, list) else ele
+    return total_sum * depth
+
+
+def product_sum_array(array: list[int | list]) -> int:
+    """
+    Calculates the product sum of an array.
+
+    Args:
+        array (List[Union[int, List]]): The array of integers and nested lists.
+
+    Returns:
+        int: The product sum of the array.
+
+    Examples:
+        >>> product_sum_array([1, 2, 3])
+        6
+        >>> product_sum_array([1, [2, 3]])
+        11
+        >>> product_sum_array([1, [2, [3, 4]]])
+        47
+        >>> product_sum_array([0])
+        0
+        >>> product_sum_array([-3.5, [1, [0.5]]])
+        1.5
+        >>> product_sum_array([1, -2])
+        -1
+
+    """
+    return product_sum(array, 1)
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()