Create a Diagonal Matrix with specified number of rows and columns in Scala

问题

I have an input mllib Block matrix named matrix like,

matrix : org.apache.spark.mllib.linalg.Matrix =
0.0  2.0  1.0  2.0
2.0  0.0  2.0  4.0
1.0  2.0  0.0  3.0
2.0  4.0  3.0  0.0

As per my Scala code, diagonals will be zero for sure. I need the diagonals of the matrix to be 1. If I have a diagonal matrix with diagonal values as 1 like,

diagonalMatrix: org.apache.spark.mllib.linalg.Matrix =
1.0  0.0  0.0  0.0
0.0  1.0  0.0  0.0
0.0  0.0  1.0  0.0
0.0  0.0  0.0  1.0

I can add those matrices, So the diagonals of matrix will be changed to 1.

matrix : org.apache.spark.mllib.linalg.Matrix =
    1.0  2.0  1.0  2.0
    2.0  1.0  2.0  4.0
    1.0  2.0  1.0  3.0
    2.0  4.0  3.0  1.0

We can create a diagonal matrix with specified number of rows and columns and diagonals as 1 based on the answer given below. But as the number of rows and columns were too big, I need an optimized solution. Or is there any better solution to make diagonals of matrix to 1 ?

回答1:

val nR = 5
val nC = 5

val seq = for {
  i <- 0 until nC
  j <- 0 until nR
  v = if (i == j) 1d else 0d
} yield v

val matrix = DenseMatrix(nR, nC, seq.toArray)

来源：https://stackoverflow.com/questions/51242763/create-a-diagonal-matrix-with-specified-number-of-rows-and-columns-in-scala