如何在R语言中按ID与分组条件计算新列？示例详解

原始数据 #

ID <- c(1,2,3,1,2,3)
t <- c(1,1,1,2,2,2)
number <- c(0.5,0.4,0.1,0.3,0.25,0.45)
dat <- cbind(ID,t,number)

原始数据输出：

ID t number
[1,]  1 1   0.50
[2,]  2 1   0.40
[3,]  3 1   0.10
[4,]  1 2   0.30
[5,]  2 2   0.25
[6,]  3 2   0.45

需求说明 #

新增calc列，规则如下：

• 当ID不等于1时，calc = 当前行number × 同t分组下ID=1对应的number值
• 当ID等于1时，calc值为NA

期望输出结果：

ID t number  calc
[1,]  1 1   0.50    NA
[2,]  2 1   0.40 0.200
[3,]  3 1   0.10 0.050
[4,]  1 2   0.30    NA
[5,]  2 2   0.25 0.075
[6,]  3 2   0.45 0.135

实现方法 #

方法一：Base R 实现 #

先将矩阵转为数据框（矩阵操作灵活性较差），再通过分组匹配计算：

# 转换为数据框
dat_df <- as.data.frame(dat)

# 提取每个t分组对应的ID=1的number值，生成映射表
t_ref <- aggregate(number ~ t, data = dat_df, subset = ID == 1, FUN = unique)

# 计算calc列
dat_df$calc <- ifelse(dat_df$ID == 1, NA, dat_df$number * t_ref$number[match(dat_df$t, t_ref$t)])

# 若需保留矩阵格式，可再转换回去
dat <- as.matrix(dat_df)

方法二：dplyr（tidyverse）实现 #

用分组函数group_by配合mutate，代码更直观：

library(dplyr)

dat_df <- as.data.frame(dat) %>%
  group_by(t) %>%
  # 每组内提取ID=1的number值，完成计算
  mutate(calc = ifelse(ID == 1, NA, number * number[ID == 1])) %>%
  ungroup()

# 如需矩阵格式，转换即可
dat <- as.matrix(dat_df)

方法三：Base R 简化写法（用ave函数） #

dat_df <- as.data.frame(dat)
dat_df$calc <- ifelse(dat_df$ID == 1, NA, 
                      dat_df$number * ave(dat_df$number, dat_df$t, 
                                         FUN = function(x) x[dat_df$ID[dat_df$t == unique(dat_df$t)] == 1]))
dat <- as.matrix(dat_df)

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