diff --git a/DIRECTORY.md b/DIRECTORY.md index 6ec8d5111..83389dab1 100644 --- a/DIRECTORY.md +++ b/DIRECTORY.md @@ -166,6 +166,7 @@ * Arrays * [Permutations](data_structures/arrays/permutations.py) * [Prefix Sum](data_structures/arrays/prefix_sum.py) + * [Product Sum Array](data_structures/arrays/product_sum.py) * Binary Tree * [Avl Tree](data_structures/binary_tree/avl_tree.py) * [Basic Binary Tree](data_structures/binary_tree/basic_binary_tree.py) diff --git a/data_structures/arrays/product_sum.py b/data_structures/arrays/product_sum.py new file mode 100644 index 000000000..4fb906f36 --- /dev/null +++ b/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()