如何基于数据框两列值作为矩阵索引提取对应值生成新列

问题背景 #

现有包含var1和var2两列的数据框df：

df = data.frame(var1 = c(3,7,2,6,14,9,3,2,1,5,7,9,5,13,2,8),
                var2 = c(5,7,3,1,5,8,3,6,5,1,4,7,5,7,8,1))

数据展示：

var1 var2
    3    5
    7    7
    2    3
    6    1
    14   5
    9    8
    3    3
    2    6
    1    5
    5    1 
    7    4
    9    7
    5    5
    13   7
    2    8 
    8    1

同时存在8行15列的矩阵x：

x = matrix(1:10, nrow = 8, ncol = 15)

矩阵内容：

[,1] [,2] [,3] [,4] [,5] [,6] [,7] [,8] [,9] [,10] [,11] [,12] [,13] [,14] 
[1,]    1    9    7    5    3    1    9    7    5     3     1     9     7     5     
[2,]    2   10    8    6    4    2   10    8    6     4     2    10     8     6     
[3,]    3    1    9    7    5    3    1    9    7     5     3     1     9     7     
[4,]    4    2   10    8    6    4    2   10    8     6     4     2    10     8     
[5,]    5    3    1    9    7    5    3    1    9     7     5     3     1     9     
[6,]    6    4    2   10    8    6    4    2   10     8     6     4     2    10     
[7,]    7    5    3    1    9    7    5    3    1     9     7     5     3     1     
[8,]    8    6    4    2   10    8    6    4    2    10     8     6     4     2    

需求：以var2为矩阵行索引、var1为矩阵列索引，提取对应位置的数值，为df添加新列var3，最终得到如下数据框：

var1 var2 var3
    3    5   1
    7    7   5
    2    3   1
    6    1   1
    14   5   9
    9    8   2
    3    3   9
    2    6   4
    1    5   5
    5    1   3
    7    4   2
    9    7   1
    5    5   7
    13   7   3
    2    8   6
    8    1   7

解决方案 #

在R中可以直接利用矩阵的索引特性，通过cbind构造索引矩阵，一次性提取所有对应元素：

# 为df添加var3列
df$var3 <- x[cbind(df$var2, df$var1)]

# 查看结果
print(df)

原理说明 #

• cbind(df$var2, df$var1)生成一个2列的矩阵，每一行对应一组行、列索引对。
• R的矩阵支持传入这种索引矩阵，会按每一行的索引提取对应位置的元素，最终返回一个与索引行数一致的向量，正好可以作为df的新列。

内容的提问来源于stack exchange，提问作者Bartholomew