R code golf: the identity matrix

How many different ways are there to create an identity matrix in R? This was an interesting little challenge set by Guillaume Nicoulaud on Twitter.

In code golf, programmers try to write an algorithm using the shortest programme possible, often exploiting lesser-known eccentricies of each programming language.

This R challenge is less about minimising your golf handicap, however, and more about showing off some features and functions in R that others might not know about. Taking the scenic route and learning along the way. Crazy code golf, if you like.

The identity matrix

As a quick reminder, the identity matrix is the linear algebraic equivalent of the number 1. It is a diagonal matrix of ones, with all off-diagonal entries equal to zero. The three-dimensional identity matrix, for example, is $$\mathbf{I} = \begin{bmatrix} 1 & 0 & 0 \\ 0 & 1 & 0 \\ 0 & 0 & 1 \end{bmatrix}.$$ Multiply any matrix $\mathbf{M}$ by the identity and you get the same matrix back. $$\mathbf{M} \mathbf{I} = \mathbf{M}$$ Multiply any matrix by its inverse $\mathbf{M}^{-1}$ (if it exists) and you get the identity matrix back. $$\mathbf{M} \mathbf{M}^{-1} = \mathbf{M}^{-1} \mathbf{M} = \mathbf{I}$$

Identity matrices in R

I will now show you a variety of ways to create a five-dimensional identity matrix in R, from the straightforward to the more esoteric.

The easy way

There is a function designed for exactly this purpose, so use it.

n <- 5L
diag(n)
     [,1] [,2] [,3] [,4] [,5]
[1,]    1    0    0    0    0
[2,]    0    1    0    0    0
[3,]    0    0    1    0    0
[4,]    0    0    0    1    0
[5,]    0    0    0    0    1

The diag() function, when passed a matrix, extracts the diagonal elements from that matrix. When passed a vector, it creates a diagonal matrix with entries equal to that vector. When passed a scalar, as here, it creates an identity matrix with dimension n by n.

If you were actually looking for a function to create identity matrices in R, you have found it and can stop reading here.

The optimistic mathematician’s way

In the equation above, we saw that the identity matrix is equal to any matrix multiplied by its own inverse. So let’s do that, using a matrix of random normal samples.

x <- matrix(rnorm(n * n), n, n)
x %*% solve(x)
              [,1]          [,2]          [,3]          [,4]          [,5]
[1,]  1.000000e+00 -1.387779e-16 -1.196959e-16 -2.498002e-16 -3.330669e-16
[2,]  0.000000e+00  1.000000e+00  2.775558e-17  0.000000e+00  2.220446e-16
[3,]  0.000000e+00  0.000000e+00  1.000000e+00  0.000000e+00  0.000000e+00
[4,] -2.220446e-16  5.551115e-17 -6.938894e-18  1.000000e+00  0.000000e+00
[5,] -5.551115e-16  1.387779e-16 -3.816392e-17  0.000000e+00  1.000000e+00

Unfortunately computers don’t have infinite precision, so there is a bit of floating-point error here. The result is almost equal to an identity matrix. A bit of rounding should take care of that:

round(x %*% solve(x))
     [,1] [,2] [,3] [,4] [,5]
[1,]    1    0    0    0    0
[2,]    0    1    0    0    0
[3,]    0    0    1    0    0
[4,]    0    0    0    1    0
[5,]    0    0    0    0    1

Perfect!

The vectorised index way

Diagonal entries are those whose row and column index are equal. That is, an identity matrix is a matrix $\mathbf{D}$ whose elements are $$d_{ij} = \begin{cases} 1 & i = j, \\ 0 & i \neq j \end{cases}.$$

Or in R terms:

.col(c(n, n)) == .row(c(n, n))
      [,1]  [,2]  [,3]  [,4]  [,5]
[1,]  TRUE FALSE FALSE FALSE FALSE
[2,] FALSE  TRUE FALSE FALSE FALSE
[3,] FALSE FALSE  TRUE FALSE FALSE
[4,] FALSE FALSE FALSE  TRUE FALSE
[5,] FALSE FALSE FALSE FALSE  TRUE

Multiply this result by 1, or add 0, to convert from boolean values to binary. If you already have an $n \times n$ matrix, you can be even more concise:

x <- matrix(nrow = n, ncol = n)
1 * (col(x) == row(x))
     [,1] [,2] [,3] [,4] [,5]
[1,]    1    0    0    0    0
[2,]    0    1    0    0    0
[3,]    0    0    1    0    0
[4,]    0    0    0    1    0
[5,]    0    0    0    0    1

Just beware that the lower-level .col() and .row() functions, unlike anything else in the R universe, actually care about the difference between an integer and a float. The command .col(c(5, 5)) produces the unhelpful error,

Error in .col(c(5, 5)) : 
  a matrix-like object is required as argument to 'col'

whilst .col(c(5L, 5L)), where the dimensions are explicitly integers, gives the desired result.

The outer product way

Similar to above, but more elegant.

outer(1:n, 1:n, '==') + 0

Thanks to Florian Privé for this one. The outer product function applies every possible combination of the elements of the first two arguments and returns the results in a matrix.

In case that wasn’t clear, here is a Battleship grid:

outer(LETTERS[1:5], 1:5, paste0)
     [,1] [,2] [,3] [,4] [,5]
[1,] "A1" "A2" "A3" "A4" "A5"
[2,] "B1" "B2" "B3" "B4" "B5"
[3,] "C1" "C2" "C3" "C4" "C5"
[4,] "D1" "D2" "D3" "D4" "D5"
[5,] "E1" "E2" "E3" "E4" "E5"

The loopy index way

As above, but using loops, which in R are quite slow and should be avoided.

x <- matrix(nrow = n, ncol = n)
for (i in 1:n)
  for (j in 1:n)
    if (i == j)
      x[i, j] <- 1

The sapply way

If you want to use loops whilst convincing yourself that you aren’t using loops.

sapply(1:n, function(i) as.integer(1:n == i))

The tidy way

The tidyverse encourages us to store everything in long data frames, so let’s not make an exception here. Here is a dataset describing the positions of all the ones.

data.frame(value = 1, row_index = 1:n, column_index = 1:n)
  row_index column_index value
1         1            1     1
2         2            2     1
3         3            3     1
4         4            4     1
5         5            5     1

Now we can use the handy xtabs() function to make a two-way contingency table. The dimnames become no longer relevant so we can remove those while we are at it.

unname( xtabs(data.frame(1, 1:n, 1:n)) )

Edit: An even simpler version is

table(data.frame(1:n, 1:n))

The sparse way

An identity matrix is the same as a permutation matrix where the order of elements is not changed: $$\{1, \dots, n\} \rightarrow \{1, \dots, n\}.$$ The Matrix package has a special class, pMatrix, for sparse permutation matrices.

library(Matrix)
as(1:n, 'pMatrix')
5 x 5 sparse Matrix of class "pMatrix"
              
[1,] | . . . .
[2,] . | . . .
[3,] . . | . .
[4,] . . . | .
[5,] . . . . |

The hard way

Matrices are just reshaped vectors. The elements of an identity matrix are (either column-wise or row-wise; it doesn’t matter because it’s symmetrical) $$\{1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, \dots, 1\}$$ i.e. a one, then $n$ zeroes, another one, another $n$ zeroes and so on, finishing with the $n$th one.

We can write down such a sequence and then reshape it.

matrix(c(rep(c(1, rep(0, n)), n - 1), 1), n)

Basically unreadable, but it works!

Guillaume’s own version doesn’t fit on one line but is a bit more sensible:

x <- rep(0, n * n)
x[seq(1, n * n, n + 1)] <- 1
dim(x) <- c(n, n)

Colin Fay makes a more readable solution, taking advantage of matrix broadcasting. That is, R will repeat a vector to fit the space allocated if it isn’t already long enough.

matrix(c(1, rep(0, n)), nrow = n, ncol = n)

Because $n^2$ isn’t an exact multiple of $n+1$, R will throw a warning, which you can ignore.

     [,1] [,2] [,3] [,4] [,5]
[1,]    1    0    0    0    0
[2,]    0    1    0    0    0
[3,]    0    0    1    0    0
[4,]    0    0    0    1    0
[5,]    0    0    0    0    1
Warning message:
In matrix(c(1, rep(0, n)), nrow = n, ncol = n) :
  data length [6] is not a sub-multiple or multiple of the number of rows [5]

If this annoys you, hide it with suppressWarnings().

The silly way

To me as a statistician, all of these methods seem are bit too… deterministic. Let’s fix that.

m <- matrix(rbinom(5^2, 1, 1/5), 5)
while(any(m %*% m != m)
      || any(colSums(m) == 0)
      || any(rowSums(m) == 0))
  m[] <- rbinom(5^2, 1, 1/5)

This glorious code will repeatedly sample entries of m from the binomial distribution until the result is an identity matrix. Arguably, this is the best solution.


With that challenge solved, what other R code golf challenges would you like to see? Maybe all the different ways of tackling a common data wrangling task, or producing a certain kind of visualisation—please let me know if you have any ideas.