如何将数据集中指定的多列one-hot编码转换为因子列？

我需要把数据集中的多组独热编码列转换为单个因子列，并且希望能指定要合并的列组来完成转换。以下是示例数据和预期结果：

library(tidyverse)

tbl <- tibble(
  # 独热编码列组1
  a1_blue = c(1, 0, 0),
  a1_red = c(0, 1, 0),
  a1_green = c(0, 0, 1),
  # 独热编码列组2
  a2_square = c(1, 0, 0),
  a2_circle = c(0, 1, 0),
  a2_dot = c(0, 0, 1),
  
  a3_letters = factor(c("A", "B", "C"))
)

# 预期结果
tbl_desired <- tibble(
  a1_colors = factor(c("blue", "red", "green"), 
              levels = c("blue", "red", "green")), 
  a2_shapes = factor(c("square", "circle", "dot"), 
              levels = c("square", "circle", "dot")), 
  a3_letters = factor(c("A", "B", "C"))
)

可以借助tidyverse的工具链实现这个需求，下面是一个可复用的函数式实现：

library(tidyverse)

# 定义合并独热编码列到因子列的函数
merge_onehot_to_factor <- function(data, group_prefix, new_col_name) {
  # 筛选当前组的所有独热编码列
  group_cols <- str_subset(names(data), glue::glue("^{group_prefix}"))
  # 提取因子水平（列名中前缀后的部分）
  factor_levels <- str_remove(group_cols, glue::glue("^{group_prefix}_"))
  
  data %>%
    # 把当前组的宽表转长表
    pivot_longer(
      cols = all_of(group_cols),
      names_to = "temp_col",
      values_to = "temp_val"
    ) %>%
    # 只保留值为1的行（对应独热编码选中的类别）
    filter(temp_val == 1) %>%
    # 转换为指定水平的因子列
    mutate(
      {{new_col_name}} := factor(
        str_remove(temp_col, glue::glue("^{group_prefix}_")),
        levels = factor_levels
      )
    ) %>%
    # 移除临时列
    select(-temp_col, -temp_val) %>%
    # 合并原数据中不属于当前组的列，保持原行顺序
    bind_cols(data %>% select(-all_of(group_cols))) %>%
    # 调整列顺序，把新生成的因子列放到前面
    select(all_of(new_col_name), everything())
}

# 对示例数据应用函数
tbl_processed <- tbl %>%
  merge_onehot_to_factor("a1", "a1_colors") %>%
  merge_onehot_to_factor("a2", "a2_shapes")

# 验证结果是否与预期一致
all.equal(tbl_processed, tbl_desired)

关键说明 #

• 函数通过group_prefix指定要合并的独热编码列组（比如a1对应所有以a1_开头的列），new_col_name指定转换后的因子列名。
• 自动提取原列名中的类别作为因子水平，保证顺序与原独热编码列的顺序一致。
• 支持多次调用，批量处理多组独热编码列。

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