Given a nested list (where sublists are of equal length), write a Python program to find the column-wise sum of the given list and return it in a new list.
Examples:
Input : [[1, 5, 3], [2, 7, 8], [4, 6, 9]] Output : [7, 18, 20] Input : [[20, 5], [2, 54], [45, 9], [72, 3]] Output : [139, 71]
Method #1:
Python3
# Python3 program to Column wise sum of nested list def column_sum(lst): res = [] for i in range ( 0 , len (lst)): s = 0 for j in range ( 0 , len (lst[i])): s + = lst[j][i] res.append(s) return res # Driver code lst = [[ 1 , 5 , 3 ], [ 2 , 7 , 8 ], [ 4 , 6 , 9 ]] print (column_sum(lst)) |
[7, 18, 20]
Time Complexity: O(n^2), where n is the length of the longest sublist in the nested list.
Auxiliary Space: O(n), where n is the number of columns in the nested list.
Method #2: zip using list comprehension We can find sum of each column of the given nested list using zip function of python enclosing it within list comprehension.
Python3
# Python3 program to Column wise sum of nested list def column_sum(lst): return [ sum (i) for i in zip ( * lst)] # Driver code lst = [[ 1 , 5 , 3 ], [ 2 , 7 , 8 ], [ 4 , 6 , 9 ]] print (column_sum(lst)) |
[7, 18, 20]
Time Complexity: O(n) where n is the number of elements in the list
Auxiliary Space: O(n), where n is the length of the list
Method #3: Using map() method Another approach is to use map(). We apply the sum function to each element in a column and find sum of each column accordingly.
Python3
# Python3 program to Column wise sum of nested list def column_sum(lst): return list ( map ( sum , zip ( * lst))) # Driver code lst = [[ 1 , 5 , 3 ], [ 2 , 7 , 8 ], [ 4 , 6 , 9 ]] print (column_sum(lst)) |
[7, 18, 20]
Time Complexity: O(n), where n is the number of elements in the list
Auxiliary Space: O(n), where n is the number of elements in the list
Method #4: Using numpy.sum() numpy.sum() function returns the sum of array elements over the specified axis.
Python3
# Python3 program to Column wise sum of nested list from numpy import array def column_sum(lst): arr = array(lst) return sum (arr, 0 ).tolist() # Driver code lst = [[ 1 , 5 , 3 ], [ 2 , 7 , 8 ], [ 4 , 6 , 9 ]] print (column_sum(lst)) |
[7, 18, 20]
Method #5: Using dictionary
Here is another approach using a dictionary to store the sums of each column. This method iterates through each element in the nested list and adds it to the corresponding key in the dictionary.
Python3
def column_sum(lst): # Initialize dictionary to store sums column_sums = {i: 0 for i in range ( len (lst[ 0 ]))} # Iterate through each element in the list for sublist in lst: for i, val in enumerate (sublist): column_sums[i] + = val # Return the values of the dictionary as a list return list (column_sums.values()) # Test with example input lst = [[ 1 , 5 , 3 ], [ 2 , 7 , 8 ], [ 4 , 6 , 9 ]] print (column_sum(lst)) #This code is contributed by Edula Vinay Kumar Reddy |
[7, 18, 20]
Time complexity: O(nm), where n is the number of sublists and m is the number of elements in each sublist.
Auxiliary Space: O(m), as we store a sum for each column in the dictionary.