Sum of main diagonal elements in a Matrix which are prime

Given a matrix mat[][] of R rows and C columns. The task is to find the sum of all elements from the main diagonal which are prime numbers. 
Note: The main diagonals are the ones that occur from Top Left of Matrix Down To Bottom Right Corner.

Examples:

Input: R = 3, C = 3, mat[][] = {{1, 2, 3}, {0, 1, 2}, {0, 4, 2}} 
Output: 2 
Explanation:
Elements from main diagonal are { 1, 1, 2}, out of these only ‘2’ is a prime number. 
Therefore, the sum of diagonal elements which are prime = 2.

Input: R = 4, C = 4, mat[][] = { {1, 2, 3, 4}, { 0, 7, 21, 12}, { 1, 2, 3, 6}, { 3, 5, 2, 31}} 
Output: 41 
Explanation:
Elements from main diagonal are { 1, 7, 3, 31}, out of these {7, 3, 31} are prime numbers. 
Therefore, the sum of diagonal elements which are prime = 7 + 3 + 31 = 41. 

Naive Approach:

1. Traverse the given matrix and check if the current element belongs to main diagonal or not.
2. If element belongs to the main diagonal and it is a prime number then add value of the element to totalSum.
3. After, traversal of the matrix print the value of totalSum.

Below is the implementation of the above approach:

C++

// C++ program for the above approach
#include <bits/stdc++.h>
using namespace std;

// Function checks whether a number
// is prime or not
bool isPrime(int n)
{
    if (n < 2) {
        return false;
    }
    // Iterate to check primality of n
    for (int i = 2; i < n; i++) {
        if (n % i == 0)
            return false;
    }

    return true;
}

// Function calculates the sum of
// prime elements of main diagonal
void primeDiagonalElementSum(
    int* mat,
    int r, int c)
{

    // Initialise total sum as 0
    int totalSum = 0;

    // Iterate the given matrix mat[][]
    for (int i = 0; i < r; i++) {

        for (int j = 0; j < c; j++) {

            int temp = *((mat + i * c) + j);

            // If element belongs to main
            // diagonal and is prime
            if ((i == j) && isPrime(temp))
                totalSum += (temp);
        }
    }

    // Print the total sum
    cout << totalSum << endl;
}

// Driver Code
int main()
{
    int R = 4, C = 5;

    // Given Matrix
    int mat[4][5] = { { 1, 2, 3, 4, 2 },
                      { 0, 3, 2, 3, 9 },
                      { 0, 4, 1, 2, 8 },
                      { 1, 2, 3, 6, 6 } };

    // Function Call
    primeDiagonalElementSum((int*)mat, R, C);
    return 0;
}

Java

// Java program for the above approach
class GFG{

// Function checks whether a number
// is prime or not
static boolean isPrime(int n)
{
    if (n < 2)
    {
        return false;
    }

    // Iterate to check primarility of n
    for(int i = 2; i < n; i++)
    {
       if (n % i == 0)
           return false;
    }

    return true;
}

// Function calculates the sum of
// prime elements of main diagonal
static void primeDiagonalElementSum(int [][]mat,
                                    int r, int c)
{

    // Initialise total sum as 0
    int totalSum = 0;

    // Iterate the given matrix mat[][]
    for(int i = 0; i < r; i++)
    {
       for(int j = 0; j < c; j++)
       {
          int temp = mat[i][j];

          // If element belongs to main
          // diagonal and is prime
          if ((i == j) && isPrime(temp))
              totalSum += (temp);
       }
    }

    // Print the total sum
    System.out.print(totalSum + "\n");
}

// Driver Code
public static void main(String[] args)
{
    int R = 4, C = 5;

    // Given Matrix
    int mat[][] = { { 1, 2, 3, 4, 2 },
                    { 0, 3, 2, 3, 9 },
                    { 0, 4, 1, 2, 8 },
                    { 1, 2, 3, 6, 6 } };

    // Function Call
    primeDiagonalElementSum(mat, R, C);
}
}

// This code is contributed by gauravrajput1

Python3

# Python3 program for the above approach

# Function checks whether a number
# is prime or not
def isPrime(n):

    if (n < 2):
        return False

    # Iterate to check primarility of n
    for i in range(2, n):
        if (n % i == 0):
            return False

    return True

# Function calculates the sum of
# prime elements of main diagonal
def primeDiagonalElementSum(mat, r, c):

    # Initialise total sum as 0
    totalSum = 0;

    # Iterate the given matrix mat[][]
    for i in range(r):
        for j in range(c):
            temp = mat[i][j]

            # If element belongs to main
            # diagonal and is prime
            if ((i == j) and isPrime(temp)):
                totalSum += (temp)

    # Print the total sum
    print(totalSum)

# Driver code
if __name__=="__main__":

    R = 4
    C = 5

    # Given Matrix
    mat = [ [ 1, 2, 3, 4, 2 ],
            [ 0, 3, 2, 3, 9 ],
            [ 0, 4, 1, 2, 8 ],
            [ 1, 2, 3, 6, 6 ] ]

    # Function call
    primeDiagonalElementSum(mat, R, C)

# This code is contributed by rutvik_56

C#

// C# program for the above approach
using System;

class GFG{

// Function checks whether a number
// is prime or not
public static bool isPrime(int n)
{
    if (n < 2)
    {
        return false;
    }

    // Iterate to check primarility of n
    for(int i = 2; i < n; i++)
    {
       if (n % i == 0)
           return false;
    }
    return true;
}

// Function calculates the sum of
// prime elements of main diagonal
public static void primeDiagonalElementSum(int [,]mat,
                                           int r, int c)
{

    // Initialise total sum as 0
    int totalSum = 0;

    // Iterate the given matrix mat[][]
    for(int i = 0; i < r; i++)
    {
       for(int j = 0; j < c; j++)
       {
          int temp = mat[i, j];

          // If element belongs to main
          // diagonal and is prime
          if ((i == j) && isPrime(temp))
              totalSum += (temp);
       }
    }

    // Print the total sum
    Console.WriteLine(totalSum);
}

// Driver Code
public static void Main(String[] args)
{
    int R = 4, C = 5;

    // Given Matrix
    int [,]mat = { { 1, 2, 3, 4, 2 },
                   { 0, 3, 2, 3, 9 },
                   { 0, 4, 1, 2, 8 },
                   { 1, 2, 3, 6, 6 } };

    // Function Call
    primeDiagonalElementSum(mat, R, C);
}
}

// This code is contributed by SoumikMondal

Javascript

<script>

// Javascript program for
// the above approach


// Function checks whether a number
// is prime or not
function isPrime( n)
{
    if (n < 2)
    {
        return false;
    }

    // Iterate to check primarility of n
    for(let i = 2; i < n; i++)
    {
    if (n % i == 0)
        return false;
    }

    return true;
}

// Function calculates the sum of
// prime elements of main diagonal
function primeDiagonalElementSum(mat,r,c)
{

    // Initialise total sum as 0
    let totalSum = 0;

    // Iterate the given matrix mat[][]
    for(let i = 0; i < r; i++)
    {
    for(let j = 0; j < c; j++)
    {
        let temp = mat[i][j];

        // If element belongs to main
        // diagonal and is prime
        if ((i == j) && isPrime(temp))
            totalSum += (temp);
    }
    }

    // Print the total sum
    document.write(totalSum + "<br>");
}

// Driver Code

    let R = 4, C = 5;

    // Given Matrix
    let mat = [[ 1, 2, 3, 4, 2 ],
               [ 0, 3, 2, 3, 9 ],
                  [ 0, 4, 1, 2, 8 ],
               [ 1, 2, 3, 6, 6 ]];

    // Function Call
    primeDiagonalElementSum(mat, R, C);


// This code is contributed by sravan kumar

</script>

Output:

3

Time Complexity: O(R*C*K), where K is the maximum element in the matrix. 
Auxiliary Space: O(1)

Efficient Approach: We can optimize the naive approach by optimizing the primality test of the number. Below are the steps for optimizing the primality test:

1. Instead of checking till N, we can check till sqrt(N) as the larger factor of N must be a multiple of smaller factor that has been already checked.
2. The algorithm can be improved further by observing that all primes are of the form 6k ± 1, with the exception of 2 and 3. This is because all integers can be expressed as (6k + i) for some integer k and for i = -1, 0, 1, 2, 3, or 4.
3. As 2 divides (6k + 0), (6k + 2), (6k + 4); and 3 divides (6k + 3). So a more efficient method is to test if N is divisible by 2 or 3, then to check through all the numbers of form 6k ± 1.

Below is the implementation of the above approach:

C++

// C++ program for the above approach
#include <bits/stdc++.h>
using namespace std;

// Function checks whether a number
// is prime or not
bool isPrime(int n)
{
    // Corner cases
    if (n <= 1)
        return false;
    if (n <= 3)
        return true;

    // This is checked so that we can skip
    // middle five numbers in below loop
    if (n % 2 == 0 || n % 3 == 0)
        return false;

    for (int i = 5; i * i <= n; i = i + 6)
        if (n % i == 0 || n % (i + 2) == 0)
            return false;

    return true;
}

// Function calculates the sum of
// prime elements of main diagonal
void primeDiagonalElementSum(
    int* mat,
    int r, int c)
{

    // Initialise total sum as 0
    int totalSum = 0;

    // Iterate the given matrix mat[][]
    for (int i = 0; i < r; i++) {

        for (int j = 0; j < c; j++) {

            int temp = *((mat + i * c) + j);

            // If element belongs to main
            // diagonal and is prime
            if ((i == j) && isPrime(temp))
                totalSum += (temp);
        }
    }

    // Print the total sum
    cout << totalSum << endl;
}

// Driver Code
int main()
{
    int R = 4, C = 5;

    // Given Matrix
    int mat[4][5] = { { 1, 2, 3, 4, 2 },
                      { 0, 3, 2, 3, 9 },
                      { 0, 4, 1, 2, 8 },
                      { 1, 2, 3, 6, 6 } };

    // Function Call
    primeDiagonalElementSum((int*)mat, R, C);
    return 0;
}

Java

// Java program for the above approach
import java.util.*;

class GFG{

// Function checks whether a number
// is prime or not
static boolean isPrime(int n)
{

    // Corner cases
    if (n <= 1)
        return false;
    if (n <= 3)
        return true;

    // This is checked so that we can skip
    // middle five numbers in below loop
    if (n % 2 == 0 || n % 3 == 0)
        return false;

    for(int i = 5; i * i <= n; i = i + 6)
        if (n % i == 0 || n % (i + 2) == 0)
            return false;

    return true;
}

// Function calculates the sum of
// prime elements of main diagonal
static void primeDiagonalElementSum(int[][] mat,
                                    int r, int c)
{

    // Initialise total sum as 0
    int totalSum = 0;

    // Iterate the given matrix mat[][]
    for(int i = 0; i < r; i++)
    {
        for(int j = 0; j < c; j++)
        {
            int temp = mat[i][j];

            // If element belongs to main
            // diagonal and is prime
            if ((i == j) && isPrime(temp))
                totalSum += (temp);
        }
    }

    // Print the total sum
    System.out.print(totalSum + "\n");
}

// Driver Code
public static void main(String[] args)
{
    int R = 4, C = 5;

    // Given Matrix
    int mat[][] = { { 1, 2, 3, 4, 2 },
                    { 0, 3, 2, 3, 9 },
                    { 0, 4, 1, 2, 8 },
                    { 1, 2, 3, 6, 6 } };

    // Function call
    primeDiagonalElementSum(mat, R, C);
}
}

// This code is contributed by Rajput-Ji

Python3

# Python3 program for the above approach

# Function checks whether a number
# is prime or not
def isPrime(n):

    # Corner cases
    if (n <= 1):
        return False
    if (n <= 3):
        return True

    # This is checked so that we can skip
    # middle five numbers in below loop
    if (n % 2 == 0 or n % 3 == 0):
        return False

    i = 5
    while i * i <= n:
        if (n % i == 0 or n % (i + 2) == 0):
            return False

        i += 6

    return True

# Function calculates the sum of
# prime elements of main diagonal
def primeDiagonalElementSum(mat, r, c):

    # Initialise total sum as 0
    totalSum = 0

    # Iterate the given matrix mat[][]
    for i in range(r):
        for j in range(c):
            temp = mat[i][j]

            # If element belongs to main
            # diagonal and is prime
            if ((i == j) and isPrime(temp)):
                totalSum += (temp)

    # Print the total sum
    print(totalSum)

# Driver Code
R, C = 4, 5

# Given Matrix
mat = [ [ 1, 2, 3, 4, 2 ],
        [ 0, 3, 2, 3, 9 ],
        [ 0, 4, 1, 2, 8 ],
        [ 1, 2, 3, 6, 6 ] ]

# Function Call
primeDiagonalElementSum(mat, R, C)

# This code is contributed by divyeshrabadiya07

C#

// C# program for the above approach
using System;
class GFG{

// Function checks whether a number
// is prime or not
static bool isPrime(int n)
{

    // Corner cases
    if (n <= 1)
        return false;
    if (n <= 3)
        return true;

    // This is checked so that we can skip
    // middle five numbers in below loop
    if (n % 2 == 0 || n % 3 == 0)
        return false;

    for(int i = 5; i * i <= n; i = i + 6)
        if (n % i == 0 || n % (i + 2) == 0)
            return false;

    return true;
}

// Function calculates the sum of
// prime elements of main diagonal
static void primeDiagonalElementSum(int[,] mat,
                                    int r, int c)
{

    // Initialise total sum as 0
    int totalSum = 0;

    // Iterate the given matrix [,]mat
    for(int i = 0; i < r; i++)
    {
        for(int j = 0; j < c; j++)
        {
            int temp = mat[i,j];

            // If element belongs to main
            // diagonal and is prime
            if ((i == j) && isPrime(temp))
                totalSum += (temp);
        }
    }

    // Print the total sum
    Console.Write(totalSum + "\n");
}

// Driver Code
public static void Main(String[] args)
{
    int R = 4, C = 5;

    // Given Matrix
    int [,]mat = { { 1, 2, 3, 4, 2 },
                    { 0, 3, 2, 3, 9 },
                    { 0, 4, 1, 2, 8 },
                    { 1, 2, 3, 6, 6 } };

    // Function call
    primeDiagonalElementSum(mat, R, C);
}
}

// This code is contributed by Rohit_ranjan

Javascript

<script>

// Javascript program for the above approach

// Function checks whether a number
// is prime or not
function isPrime(n)
{
    // Corner cases
    if (n <= 1)
        return false;
    if (n <= 3)
        return true;

    // This is checked so that we can skip
    // middle five numbers in below loop
    if (n % 2 == 0 || n % 3 == 0)
        return false;

    for (var i = 5; i * i <= n; i = i + 6)
        if (n % i == 0 || n % (i + 2) == 0)
            return false;

    return true;
}

// Function calculates the sum of
// prime elements of main diagonal
function primeDiagonalElementSum(  mat, r, c)
{

    // Initialise total sum as 0
    var totalSum = 0;

    // Iterate the given matrix mat[][]
    for (var i = 0; i < r; i++) {

        for (var j = 0; j < c; j++) {

            var temp = mat[i][j];

            // If element belongs to main
            // diagonal and is prime
            if ((i == j) && isPrime(temp))
                totalSum += (temp);
        }
    }

    // Print the total sum
    document.write( totalSum );
}

// Driver Code
var R = 4, C = 5;
// Given Matrix
var mat = [ [ 1, 2, 3, 4, 2 ],
                  [ 0, 3, 2, 3, 9 ],
                  [ 0, 4, 1, 2, 8 ],
                  [ 1, 2, 3, 6, 6 ] ];
// Function Call
primeDiagonalElementSum(mat, R, C);

// This code is contributed by itsok.
</script>

Output:

3

Time Complexity: O(R*C*sqrt(K)), where K is the maximum element in the matrix. 
Auxiliary Space: O(1)