replace array values according to values in the first index in the first dimension

Question

I have an array with a few dimensions. I want to replace values according to values in the first index in the first dimension. In the example below, I want to change all values that the corresponding a1 dimension == 2. If I change only one index:

set.seed(2)
arr <- array(data=sample(1:2, 18, replace = TRUE), dim=c(3,3,2), dimnames=list(paste0("a",1:3),paste0("b",1:3),paste0("c",1:2)))
    
# replace second index according to first index of dimension 1
arr[2,,][arr[1,,]==2] <- NA

The result is as expected:

> arr
, , c1

   b1 b2 b3
a1  1  2  1
a2  1 NA  1
a3  2  2  1

, , c2

   b1 b2 b3
a1  2  2  1
a2 NA NA  2
a3  1  1  2

But if I try to change all other indexes like this:

set.seed(2)
    arr <- array(data=sample(1:2, 18, replace = TRUE), dim=c(3,3,2), dimnames=list(paste0("a",1:3),paste0("b",1:3),paste0("c",1:2)))
        
        # replace 2nd & 3rd index according to first index of dimension 1
        arr[2:3,,][arr[1,,]==2] <- NA

It doesn't work as I expect. The indexes in an array is difficult to understand. How do I do it correctly? (naturally, without changing each index separately). Thanks.

I expect the result to be:

> arr
, , c1

   b1 b2 b3
a1  1  2  1
a2  1 NA  1
a3  2 NA  1

, , c2

   b1 b2 b3
a1  2  2  1
a2 NA NA  2
a3 NA NA  2

Answer 1

You can use rep to get the right indices for subsetting.

arr[2:3,,][rep(arr[1,,]==2, each=2)] <- NA

arr
#, , c1
#
#   b1 b2 b3
#a1  1  2  1
#a2  1 NA  1
#a3  2 NA  1
#
#, , c2
#
#   b1 b2 b3
#a1  2  2  1
#a2 NA NA  2
#a3 NA NA  2

Or more generally.

i <- 2:dim(arr)[1]
arr[i,,][rep(arr[1,,]==2, each=length(i))] <- NA

Or (Thanks to @jblood94 for this variant)

arr[-1,,][rep(arr[1,,]==2, each = nrow(arr) - 1)] <- NA

Or using a loop.

for(i in 2:nrow(arr)) arr[i,,][arr[1,,]==2] <- NA

Answer 2

It would be

arr[2:3,,][rep(arr[1,,]==2, each = 2)] <- NA

Or, more generally, to replace all rows based on the first row:

arr[-1,,][rep(arr[1,,]==2, each = nrow(arr) - 1)] <- NA