#include <iostream>

#include <vector>

class Matrix {
private:
std::vector<std::vector<double>> data;

public:
// Constructors and other methods...

// Method to count and return the number of non-zero elements

int nnz() const {
int count = 0;
for (const auto& row : data) {
for (double element : row) {
if (element != 0) {
count++;
}
}
}
return count;
}

// Friend function to create a diagonal matrix or extract diagonal

elements
friend Matrix diag(const Matrix& mv);
};

// Implementation of the external function diag

Matrix diag(const Matrix& mv) {
int rows = mv.data.size();
int cols = mv.data[0].size();

Matrix result;

if (rows == cols) {
// Square matrix, create a column vector of the main diagonal
elements
result.data.resize(rows, std::vector<double>(1));
for (int i = 0; i < rows; ++i) {
result.data[i][0] = mv.data[i][i];
}
} else if (cols == 1) {
 // Column vector, create a square diagonal matrix
result.data.resize(rows, std::vector<double>(rows, 0.0));
for (int i = 0; i < rows; ++i) {
result.data[i][i] = mv.data[i][0];
}
} else {
// Non-square matrix or row vector, return an empty matrix
std::cerr << "Error: Unsupported matrix for diag function." <<
std::endl;
}

return result;
}

int main() {
// Test cases
Matrix squareMatrix{{1, 2, 3}, {4, 5, 6}, {7, 8, 9}};
Matrix columnVector{{1}, {2}, {3}};
Matrix rowVector{{1, 2, 3}};

Matrix result1 = diag(squareMatrix);

Matrix result2 = diag(columnVector);
Matrix result3 = diag(rowVector);

std::cout << "Result 1 (Diagonal elements of square matrix):\n";

result1.print(); // Assume there's a print() method in your Matrix
class

std::cout << "\nResult 2 (Diagonal matrix from column vector):\n";

result2.print();

std::cout << "\nResult 3 (Error - unsupported matrix):\n";

result3.print();

// Count non-zero elements in squareMatrix

int nonZeroCount = squareMatrix.nnz();
std::cout << "\nNumber of non-zero elements in squareMatrix: "
<< nonZeroCount << std::endl;

return 0;
}