简体   繁体   中英

How to find the sum of elements above and below the diagonal of a matrix in python?

I need to find the sum of the elements that are above and below the main diagonal. I have no idea how to condition the algorithm to sum only those numbers. This is the code I have so far, with A being the matrix

A = []
N = int(raw_input("Input matrix size: "))
for i in range(0, N):
    row = []
    for j in range(0, N):
        row.append(int(raw_input("Input elements: ")))
    A.append(row)
sum = 0
for i in range(0, N):
    sum += A[i][i]
print sum       
sum2 = 0
for i in range(0, N):
    for j in range(i+1, N):
        sum2 += A[i][j]
print sum2

I am guessing I should use more for statements. Thank you

Here is an example case demonstrating how to find the sum in both cases, using nested loop:

matrix = [[i+j for j in range(4)] for i in range(4)]

for row in matrix:
    print(" ".join(list(map(str,row))))

totalSum = 0
for i in range(1,len(matrix)):
    for j in range(i):
        totalSum += matrix[i][j]
print("Below sum: ", totalSum)

totalSum = 0
for i in range(len(matrix)):
    for j in range(i+1,len(matrix)):
        totalSum += matrix[i][j]
print("Above sum: ", totalSum)

Output:

0 1 2 3
1 2 3 4
2 3 4 5
3 4 5 6
Below sum:  18
Above sum:  18

You can also use these one-liners:

Below diagonal:

totalSum = sum([matrix[i][j] for i in range(1,len(matrix)) for j in range(i)])

Above diagonal:

totalSum = sum([matrix[i][j] for i in range(len(matrix)) for j in range(i+1,len(matrix))])

If you want to sum all numbers above and below the main diagonal you can perform an index check:

totalSum = 0
for i in range(len(matrix)):
    for j in range(len(matrix)):
        if not i==j:
            totalSum += matrix[i][j]
print("Sum: ", totalSum)

Yet, another way of finding that sum (although not recommended) is to find the total sum of the matrix and the sum of the main diagonal and then perform a subtraction to find the final sum:

matrix = [[i+j for j in range(4)] for i in range(4)]

for row in matrix:
    print(" ".join(list(map(str,row))))

matrixSum = sum([sum(elem for elem in row) for row in matrix])
diagonalSum = sum([matrix[i][i] for i in range(len(matrix))])
finalSum = matrixSum - diagonalSum

print("Matrix sum: ", matrixSum)
print("Diagonal sum: ", diagonalSum)
print("Final sum: ", finalSum)

Output:

0 1 2 3
1 2 3 4
2 3 4 5
3 4 5 6

Matrix sum:  48
Diagonal sum:  12
Final sum:  36

Note: Mind the syntax in the print statements as you're using Python 2 and my answer is in Python 3.

You can use np.triu , np.tril and np.trace to compute these sums (your question does not specify whether or not you are allowed to leverage numpy ):

import numpy as np

np.random.seed(0)
A = np.random.randint(0,10,size=(5,5))

Gives:

[[5 0 3 3 7]
 [9 3 5 2 4]
 [7 6 8 8 1]
 [6 7 7 8 1]
 [5 9 8 9 4]]

Then:

upper_sum = np.triu(A).sum()-np.trace(A)
lower_sum = np.tril(A).sum()-np.trace(A)

Yields:

34
73

Let's assume you have a 3x3 matrix.

[[a11, a12, a13],
 [a21, a22, a23],
 [a31, a32, a33]]

Do you find any similarity in the indices of the upper triangular and lower triangular parts? (Hover on the below text part to know the answer).

The first index of the upper triangle is lower in value as compared to the second index. Similarly, the first index is smaller than the second index in the lower triangle. Moreover, for diagonal elements, the indices are same!

Now, assuming you have written the above code yourself, I hope that you can do this one yourself, for now you know the logic. It will be two loops (one to traverse the rows, and one to traverse the columns), and some if statements.

You can acomplish this with numpy.triu_indices . I have commented below each step to guide you through it. Basically you get the upper right indices with numpy.triu_indices and loop over them to get the elements. You sum all of the elements except the ones in the diagonal.

import numpy as np

m = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
triu = np.triu_indices(m.ndim+1) # Get upper right indices of m
triu_idx = list(zip(triu[0],triu[1])) # List of tuples for the upper right indices
cum_sum = 0 # Initialize sum
for x in triu_idx: # Loop over the upper right indices
  if x[0] != x[1]: # Exclude diagonal elements
    cum_sum += m[x] # Add to sum

print(cum_sum) # Output 11

Given the matrix

[[1 2 3]
 [4 5 6]
 [7 8 9]]

It outputs 11.

Here is a fast method using scipy.spatial.distance.squareform for the triangles and np.einsum for the diagonal:

>>> import numpy as np
>>> from scipy.spatial.distance import squareform
>>> 
>>> x = np.arange(36).reshape(6, 6)
>>>
>>> sum_ut = squareform(x, checks=False).sum()
>>> sum_dg = np.einsum('ii', x)
>>> sum_lt = squareform(x.T, checks=False).sum()

Timings:

>>> timeit(lambda: squareform(x, checks=False).sum())
6.272806407185271
>>> timeit(lambda: np.einsum('ii', x))
1.3961836302187294
>>> timeit(lambda: squareform(x.T, checks=False).sum())
6.6827554509509355

For comparison:

>>> timeit(lambda: np.triu(x, 1).sum())
13.85556498519145
>>> timeit(lambda: np.trace(x))
3.081781509099528
>>> timeit(lambda: np.tril(x, -1).sum())
13.659938262077048

The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM