Program to find the sum of elements that forms a Z shape on matrix in Python

Suppose we have one n x n matrix M, we have to find the sum of all elements that form a Z shape in the matrix.

So, if the input is like

then the output will be 23, as elements are [4+3+2+1+2+5+6] = 23.

To solve this, we will follow these steps −

• n := row count of matrix
• if n <= 2, then
  • return sum of all elements in matrix
• first_row := sum of first row
• last_row := sum of last row
• diagonal = sum of matrix[i, n-1-i] for all i from 1 to n-2
• return first_row + last_row + diagonal

Let us see the following implementation to get better understanding −

Example

class Solution:
   def solve(self, matrix):
      n = len(matrix)
      if n <= 2:
         return sum(sum(row) for row in matrix) first_row = sum(matrix[0])
      last_row = sum(matrix[n-1])
      diagonal = sum(matrix[i][n-1-i] for i in range(1, n-1))
      return first_row + last_row + diagonal
ob = Solution()
matrix = [ [4, 3, 2], [9, 1, 8], [2, 5, 6] ]
print(ob.solve(matrix))

Input

matrix = [[4, 3, 2],
[9, 1, 8],
[2, 5, 6]]

Output

23