Problem Description

You are given a 2D integer array nums with equal rows and columns (a square matrix). Your task is to find the largest prime number that appears on either of the two main diagonals of this matrix.

The two diagonals are:

1. Primary diagonal: Elements where row index equals column index (nums[i][i])
2. Secondary diagonal: Elements where the row index i corresponds to column index nums.length - i - 1

For example, in a 3×3 matrix:

• Primary diagonal contains elements at positions [0][0], [1][1], [2][2]
• Secondary diagonal contains elements at positions [0][2], [1][1], [2][0]

A prime number is defined as an integer greater than 1 that has no positive divisors other than 1 and itself.

If no prime numbers exist on either diagonal, return 0.

The solution iterates through each row of the matrix and checks two elements:

• The element on the primary diagonal at position [i][i]
• The element on the secondary diagonal at position [i][n-i-1]

For each element, it uses the is_prime function to determine if it's prime. The is_prime function checks divisibility from 2 up to the square root of the number. The algorithm keeps track of the maximum prime number found and returns it at the end.

Intuition

The problem asks us to find the largest prime number on the diagonals, which naturally breaks down into two subproblems: identifying which elements are on the diagonals and checking if they are prime.

For the diagonal elements, we can observe a pattern. In a square matrix of size n×n, the primary diagonal elements follow the pattern where both indices are the same: [0][0], [1][1], [2][2], etc. The secondary diagonal has a complementary relationship: when we're at row i, the column index is n-i-1. This makes sense because as we go down rows (increasing i), we need to go left in columns (decreasing from n-1 to 0).

Since we need to find the maximum prime, we need to check all elements on both diagonals. We can do this efficiently in a single pass through the rows. For each row i, we check exactly two elements: nums[i][i] and nums[i][n-i-1]. Note that when n is odd, the center element nums[n//2][n//2] appears on both diagonals, but checking it twice doesn't affect our maximum calculation.

For prime checking, we use the standard optimization of only checking divisors up to sqrt(x). This works because if a number x has a divisor greater than sqrt(x), it must also have a corresponding divisor less than sqrt(x). So if we don't find any divisors up to sqrt(x), the number is prime.

The overall strategy is straightforward: traverse once through the matrix rows, check the two diagonal elements in each row for primality, and keep track of the maximum prime found. If no primes are found, the answer remains 0.

Learn more about Math patterns.

Solution Approach

The solution uses a straightforward simulation approach with mathematical optimization for prime checking.

Prime Checking Function: The is_prime function implements an optimized primality test:

• First, it handles the edge case where x < 2, returning False since numbers less than 2 are not prime
• For numbers ≥ 2, it checks if any number from 2 to sqrt(x) divides x
• The expression all(x % i for i in range(2, int(sqrt(x)) + 1)) returns True only if x is not divisible by any number in this range
• This optimization reduces time complexity from O(n) to O(√n) for each prime check

Main Algorithm: The solution iterates through the matrix once:

1. Initialize ans = 0 to store the maximum prime found
2. Get the matrix dimension n = len(nums)
3. For each row index i and corresponding row row:
   • Check the primary diagonal element row[i]:
     • If it's prime, update ans = max(ans, row[i])
   • Check the secondary diagonal element row[n - i - 1]:
     • If it's prime, update ans = max(ans, row[n - i - 1])

Time Complexity Analysis:

• We visit each row once: O(n) iterations
• For each row, we check at most 2 elements for primality
• Each prime check takes O(√m) where m is the value being checked
• Overall: O(n × √m) where m is the maximum value in the diagonals

Space Complexity:

• The solution uses O(1) extra space, only storing the answer and loop variables
• The is_prime function's range generator uses O(1) space due to lazy evaluation

The algorithm efficiently finds the answer in a single pass without storing diagonal elements separately, making it both time and space efficient.

Example Walkthrough

Let's walk through the solution with a concrete example to illustrate how the algorithm works.

Consider the following 4×4 matrix:

nums = [[11, 2,  3,  4],
        [5,  6,  7,  8],
        [9,  10, 11, 12],
        [13, 14, 15, 17]]

Step 1: Initialize variables

• ans = 0 (to track the maximum prime)
• n = 4 (matrix dimension)

Step 2: Process each row

Row 0 (i=0):

• Primary diagonal element: nums[0][0] = 11
  • Check if 11 is prime: Test divisors from 2 to √11 ≈ 3
  • 11 % 2 = 1 (not divisible)
  • 11 % 3 = 2 (not divisible)
  • 11 is prime! Update ans = max(0, 11) = 11
• Secondary diagonal element: nums[0][4-0-1] = nums[0][3] = 4
  • Check if 4 is prime: 4 % 2 = 0 (divisible by 2)
  • 4 is not prime, ans remains 11

Row 1 (i=1):

• Primary diagonal element: nums[1][1] = 6
  • Check if 6 is prime: 6 % 2 = 0 (divisible by 2)
  • 6 is not prime, ans remains 11
• Secondary diagonal element: nums[1][4-1-1] = nums[1][2] = 7
  • Check if 7 is prime: Test divisors from 2 to √7 ≈ 2
  • 7 % 2 = 1 (not divisible)
  • 7 is prime! Update ans = max(11, 7) = 11

Row 2 (i=2):

• Primary diagonal element: nums[2][2] = 11
  • Check if 11 is prime: Yes (same check as before)
  • Update ans = max(11, 11) = 11
• Secondary diagonal element: nums[2][4-2-1] = nums[2][1] = 10
  • Check if 10 is prime: 10 % 2 = 0 (divisible by 2)
  • 10 is not prime, ans remains 11

Row 3 (i=3):

• Primary diagonal element: nums[3][3] = 17
  • Check if 17 is prime: Test divisors from 2 to √17 ≈ 4
  • 17 % 2 = 1, 17 % 3 = 2, 17 % 4 = 1 (not divisible by any)
  • 17 is prime! Update ans = max(11, 17) = 17
• Secondary diagonal element: nums[3][4-3-1] = nums[3][0] = 13
  • Check if 13 is prime: Test divisors from 2 to √13 ≈ 3
  • 13 % 2 = 1, 13 % 3 = 1 (not divisible by any)
  • 13 is prime! Update ans = max(17, 13) = 17

Step 3: Return result The algorithm returns ans = 17, which is the largest prime number found on either diagonal.

Summary of diagonal elements checked:

• Primary diagonal: [11, 6, 11, 17] → primes: [11, 11, 17]
• Secondary diagonal: [4, 7, 10, 13] → primes: [7, 13]
• Maximum prime found: 17

Solution Implementation

Time and Space Complexity

Time Complexity: O(n × √M)

The algorithm iterates through each row of the n×n matrix exactly once, giving us O(n) iterations. For each row, we check at most 2 diagonal elements (primary diagonal element at position [i][i] and anti-diagonal element at position [i][n-i-1]). For each diagonal element, we call the is_prime function.

The is_prime function has a time complexity of O(√x) where x is the value being checked. In the worst case, x can be as large as the maximum value M in the matrix. The function iterates from 2 to √x to check divisibility.

Therefore, the overall time complexity is O(n) iterations × O(√M) per prime check = O(n × √M).

Space Complexity: O(1)

The algorithm uses only a constant amount of extra space:

• Variable n to store the matrix dimension
• Variable ans to track the maximum prime found
• Variable i for iteration
• Variables within the is_prime function (parameter x and loop variable i)

No additional data structures that grow with input size are used, resulting in constant space complexity O(1).

Learn more about how to find time and space complexity quickly.

Common Pitfalls

1. Counting the Center Element Twice in Odd-Sized Matrices

One of the most common mistakes is not recognizing that in odd-sized square matrices, the center element belongs to both diagonals. For example, in a 3×3 matrix, the element at position [1][1] is part of both the primary diagonal and the secondary diagonal.

The Issue: When iterating through the matrix, the code checks both nums[i][i] and nums[i][n-i-1]. For the middle row of an odd-sized matrix where i = n//2, both expressions point to the same element. While this doesn't affect correctness (we're looking for the maximum, so checking the same element twice doesn't change the result), it leads to unnecessary duplicate work.

Solution: Add a condition to skip the duplicate check:

for i in range(n):
    # Check main diagonal
    if is_prime(nums[i][i]):
        max_prime = max(max_prime, nums[i][i])
  
    # Check anti-diagonal only if it's not the same element
    if i != n - i - 1:  # Avoid checking center element twice
        if is_prime(nums[i][n - i - 1]):
            max_prime = max(max_prime, nums[i][n - i - 1])

2. Inefficient Prime Checking for Small Numbers

The is_prime function can be optimized for small numbers by handling special cases explicitly. Currently, it performs a loop even for numbers like 2 and 3, which are trivially prime.

Solution: Add early returns for small primes:

def is_prime(num: int) -> bool:
    if num < 2:
        return False
    if num == 2:
        return True
    if num % 2 == 0:  # Even numbers > 2 are not prime
        return False
  
    # Only check odd divisors starting from 3
    for divisor in range(3, int(sqrt(num)) + 1, 2):
        if num % divisor == 0:
            return False
    return True

3. Not Handling Edge Cases in Prime Checking

The current implementation handles negative numbers correctly (returns False for num < 2), but developers sometimes forget that 1 is not a prime number. Additionally, the square root calculation could theoretically cause issues with very large numbers due to floating-point precision.

Solution: Use integer arithmetic to avoid floating-point issues:

def is_prime(num: int) -> bool:
    if num < 2:
        return False
  
    divisor = 2
    while divisor * divisor <= num:  # Avoids sqrt calculation
        if num % divisor == 0:
            return False
        divisor += 1
    return True

4. Assuming Matrix is Always Non-Empty

While the problem states we have a square matrix, defensive programming suggests checking for empty matrices to avoid index errors.

Solution: Add a guard clause:

def diagonalPrime(self, nums: List[List[int]]) -> int:
    if not nums or not nums[0]:
        return 0
  
    n = len(nums)
    # ... rest of the code