Given a matrix with N rows and M columns, print all elements present on the main diagonal.
Example 1
Input: mat = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
Output: 1 5 9
Explanation: Main diagonal cells are mat[0][0], mat[1][1], and mat[2][2].
Example 2
Input: mat = [[10, 20, 30], [40, 50, 60]]
Output: 10 50
Explanation: Rectangular matrix diagonal stops after min(2, 3) cells, so only mat[0][0] and mat[1][1] get printed.
Approach
The main diagonal contains only those elements whose row and column indices are the same. So, instead of traversing the entire matrix, a single loop can directly visit these positions. The traversal stops at the smaller of the row and column counts, making it work for both square and rectangular matrices.
Algorithm
Find the number of rows and columns in the matrix, as these determine the boundaries of the main diagonal.
Consider the smaller of the row and column counts as the traversal limit, since a diagonal element requires both a valid row and column index.
Use a single loop from
0tolimit - 1, as the row and column indices of every main diagonal element are the same.Access and print
mat[index][index], which directly represents the current main diagonal element.Continue until all valid main diagonal positions have been processed.
Dry Run
diagonal print
Solution
#include <bits/stdc++.h>using namespace std;class Solution {public: // Function to print main diagonal elements of a matrix. void printMainDiagonal(vector<vector<int>>& mat) { // Store total number of rows. int n = mat.size(); // Handle an empty matrix. if (n == 0) { cout << "\n"; return; } // Store total number of columns. int m = mat[0].size(); // Store the count of valid diagonal cells. int limit = min(n, m); // Traverse every valid diagonal index. for (int index = 0; index < limit; index++) { // Print a space before every value except the first value. if (index > 0) { cout << " "; } // Print the current main diagonal element. cout << mat[index][index]; } // Move output cursor to the next line after printing finishes. cout << "\n"; }};// Driver code.int main() { vector<vector<int>> mat = {{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}; Solution sol; sol.printMainDiagonal(mat); return 0;}Complexity Analysis
Time Complexity: O(min(N, M)), traversal visits one diagonal cell for every valid index.
Space Complexity: O(1), only loop variables are used apart from output handling.
Interview follow-up Questions
Cells with equal row and column indices belong to the main diagonal.
Be the first to add a comment.