R语言按id分组批量为多列国家虚拟变量赋值1的方法

问题说明 #

• 输入数据为包含非唯一id字段、country字段的data.frame，同时存在多个以国家名称命名的虚拟变量列，所有虚拟变量初始值为0
• 赋值规则：同一id下，只要任意一行的country字段出现过某国家名称，该id对应的所有行中，匹配该国家名的虚拟变量列统一赋值为1
• 实现约束：不使用for循环、不使用pivot系列长宽表转换函数，需兼容「虚拟变量列存在但country字段无对应取值」的场景（如示例中的Belgium列）
• 初始示例数据：

df <- structure(list(id = c(1, 1, 2, 3, 4, 4), country = c("USA", "Japan", 
"Germany", "Austria", "Japan", "Austria"), USA = c(0, 0, 0, 0, 
0, 0), Germany = c(0, 0, 0, 0, 0, 0), Japan = c(0, 0, 0, 0, 0, 
0), Austria = c(0, 0, 0, 0, 0, 0), Belgium = c(0, 0, 0, 0, 0, 
0)), class = "data.frame", row.names = c(NA, -6L)) 

• 此前尝试的df[levels(factor(df$country))] = model.matrix(~country - 1, df)方案仅能逐行赋值，无法按id分组填充，错误结果如下：

structure(list(id = c(1, 1, 2, 3, 4, 4), country = c("USA", "Japan", 
"Germany", "Austria", "Japan", "Austria"), USA = c(1, 0, 0, 0, 
0, 0), Germany = c(0, 0, 1, 0, 0, 0), Japan = c(0, 1, 0, 0, 1, 
0), Austria = c(0, 0, 0, 1, 0, 1), Belgium = c(0, 0, 0, 0, 0, 
0)), row.names = c(NA, -6L), class = "data.frame") 

可用实现方案 #

方案1：dplyr 实现（基于mutate()+cur_column()逻辑） #

核心逻辑是按id分组后，遍历每个虚拟变量列，判断列名（国家名）是否在当前分组的country取值集合中，自动返回0/1值，天然兼容无对应country取值的虚拟变量列：

library(dplyr)

# 提取所有虚拟变量列名，若列顺序固定也可直接用3:ncol(df)取列
dummy_cols <- c("USA", "Germany", "Japan", "Austria", "Belgium")

df <- df %>%
  group_by(id) %>%
  mutate(across(all_of(dummy_cols), ~ as.numeric(cur_column() %in% country))) %>%
  ungroup()

方案2：base R 无依赖实现 #

先聚合得到每个id对应的国家集合，生成id维度的虚拟变量映射表，再匹配回原数据，同样兼容缺失country取值的虚拟变量列：

# 聚合得到每个id对应的所有去重国家
id_country_unique <- unique(df[, c("id", "country")])
# 生成id维度的虚拟变量矩阵
id_dummy_mat <- as.data.frame.matrix(xtabs(~ id + country, id_country_unique))
# 补全原数据中存在但country字段无取值的虚拟变量列
missing_dummy_cols <- setdiff(names(df)[3:ncol(df)], names(id_dummy_mat))
id_dummy_mat[missing_dummy_cols] <- 0
# 按id匹配回原数据表
df[, names(id_dummy_mat)] <- id_dummy_mat[match(df$id, row.names(id_dummy_mat)), ]

正确运行结果 #

两种方案运行后均可得到符合预期的结果：

structure(list(id = c(1, 1, 2, 3, 4, 4), country = c("USA", "Japan", 
"Germany", "Austria", "Japan", "Austria"), USA = c(1, 1, 0, 0, 
0, 0), Germany = c(0, 0, 1, 0, 0, 0), Japan = c(1, 1, 0, 0, 1, 
1), Austria = c(0, 0, 0, 1, 1, 1), Belgium = c(0, 0, 0, 0, 0, 
0)), class = "data.frame", row.names = c(NA, -6L))

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