如何为DataFrame批量生成基于特定行值的分类列？

需求说明 #

现有包含研究受试者编号（studynr）、访视日期（med_dt）、用药总数（TOTAL）的DataFrame dset，需要新增分类列baseline_med：

• 若受试者首次访视的TOTAL≤8，该受试者所有行赋值为A
• 若首次访视的TOTAL＞8，该受试者所有行赋值为B

示例数据 #

生成虚拟数据的R代码：

dates <- seq(c(ISOdate(2010,1,1)), c(ISOdate(2020,1,1)), by = "day")
dset <- data.frame(studynr = c(1,1,1,2,2,2,3,3,3,4,4,4,5,5,5), 
                   med_dt = (sample(dates, size=15, replace=TRUE)),
                   TOTAL = c(7,9,10,8,7,9,10,8,6,5,8,10,15,12,10))

数据输出：

studynr              med_dt TOTAL
1        1 2018-03-27 12:00:00     7
2        1 2015-03-16 12:00:00     9
3        1 2010-04-01 12:00:00    10
4        2 2018-08-17 12:00:00     8
5        2 2016-10-08 12:00:00     7
6        2 2019-06-23 12:00:00     9
7        3 2018-06-25 12:00:00    10
8        3 2017-07-09 12:00:00     8
9        3 2011-09-02 12:00:00     6
10       4 2010-07-05 12:00:00     5
11       4 2011-02-11 12:00:00     8
12       4 2012-04-14 12:00:00    10
13       5 2017-09-20 12:00:00    15
14       5 2015-08-31 12:00:00    12
15       5 2017-08-22 12:00:00    10

预期结果 #

studynr              med_dt TOTAL baseline_med
1        1 2018-03-27 12:00:00     7            B
2        1 2015-03-16 12:00:00     9            B
3        1 2010-04-01 12:00:00    10            B
4        2 2018-08-17 12:00:00     8            A
5        2 2016-10-08 12:00:00     7            A
6        2 2019-06-23 12:00:00     9            A
7        3 2018-06-25 12:00:00    10            A
8        3 2017-07-09 12:00:00     8            A
9        3 2011-09-02 12:00:00     6            A
10       4 2010-07-05 12:00:00     5            A
11       4 2011-02-11 12:00:00     8            A
12       4 2012-04-14 12:00:00    10            A
13       5 2017-09-20 12:00:00    15            B
14       5 2015-08-31 12:00:00    12            B
15       5 2017-08-22 12:00:00    10            B

当前尝试的代码 #

尝试先排序、提取首次访视数据并赋值，但无法将结果映射回原数据集：

dset <- dset[order(dset$studynr, dset$med_dt)]
dset <- data_unique(dset, select = "studynr", keep = "first")
dset$baseline_meds <- cut(dset$TOTAL,
                breaks=c(0, 9, Inf),
                labels=c('A', 'B'),  
                right = FALSE)

解决方案 #

方法1：使用dplyr（推荐，代码简洁直观） #

通过分组后排序，基于首次访视的TOTAL值生成分类列：

library(dplyr)

dset <- dset %>%
  group_by(studynr) %>%
  arrange(med_dt, .by_group = TRUE) %>%  # 按受试者分组后，对访视日期排序
  mutate(baseline_med = ifelse(first(TOTAL) <= 8, "A", "B")) %>%  # 用首次访视的TOTAL判断
  ungroup()

方法2：使用Base R #

先提取每个受试者首次访视的TOTAL，再通过合并将分类映射回原数据集：

# 按studynr和med_dt排序
dset_sorted <- dset[order(dset$studynr, dset$med_dt), ]
# 获取每个studynr的首次TOTAL
first_total <- aggregate(TOTAL ~ studynr, data = dset_sorted, FUN = head, n = 1)
# 生成分类
first_total$baseline_med <- ifelse(first_total$TOTAL <= 8, "A", "B")
# 合并回原数据集
dset <- merge(dset, first_total[, c("studynr", "baseline_med")], by = "studynr")
# 可选：恢复原数据的行顺序
dset <- dset[match(rownames(dset), rownames(dset_sorted)), ]

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