如何使用现有列中的值创建新列，以告知新值将来自哪一列？

Question

Here is an example data.

testdata <- data.frame(A = c(1,0,1,1,0,0),
                   B = c(2,0,0,0,0,1),
                   D0 = c("A","A","B","C","A","A"),
                   D1 = c("B","C","C","A","B","B"),
                   D2 = c("C", NA,NA,NA,NA,NA),
                   stringsAsFactors = F)

What I wanted to do is make a new column based on columns A and B (eg, columns Aprime and Bprime ). The values that will be placed in the new column will be from columns with D (eg, D0, D1, and D2 ). And the value in columns A and B tells which D column to pick. So for example, for the new column Aprime , the first value will be "B" because the first row of A is 1, thus it should take the first row of the D1 column. For the first row of Bprime, it should have "C" , because the first B is 2, thus it should take the first D2 value. The result should be something like this:

  A B D0 D1   D2 Aprime Bprime
1 1 2  A  B    C      B      C
2 0 0  A  C <NA>      A      A
3 1 0  B  C <NA>      C      B
4 1 0  C  A <NA>      A      C
5 0 0  A  B <NA>      A      A
6 0 1  A  B <NA>      A      B

I used the ifelse statements below to come up with the above results:

testdata$Aprime <- ifelse(testdata$A == 0, testdata$D0, ifelse(testdata$A == 1, testdata$D1, testdata$D2))
testdata$Bprime <- ifelse(testdata$B == 0, testdata$D0, ifelse(testdata$B == 1, testdata$D1, testdata$D2))

However, I would like a more generic one because the D columns are not fixed (eg, there can be D3 up to D20). How can I do this one without writing an ifelse for the Ds greater than 0 (ie., D1 and so on)?

TIA.

Answer 1

Here is a base R method using matrix subsetting to select the values and lapply to loop through columns A and B.

testdata[c("aprime", "bprime")] <-
      lapply(testdata[c("A", "B")],
             function(x) testdata[, 3:5][cbind(seq_len(nrow(testdata)), x + 1)])

The left side provides names for the new variables. On the right, the first argument of lapply provides the set of variables to run through. The second argument of lapply , testdata[, 3:5][cbind(seq_len(nrow(testdata)), x + 1)] first subsets the data.frame into the indexing columns, (D0-D2), and then provides a matrix for subsetting using cbind . The row indices are selected with seq_len..nrow and the columns are selected from the varaibles provided in the first argument of lapply .

This returns

testdata
  A B D0 D1   D2 aprime bprime
1 1 2  A  B    C      B      C
2 0 0  A  C <NA>      A      A
3 1 0  B  C <NA>      C      B
4 1 0  C  A <NA>      A      C
5 0 0  A  B <NA>      A      A
6 0 1  A  B <NA>      A      B

For more information on matrix subsetting, take a look at ?"[" .