R语言基于多列关联条件为数据框n4列填充red值的方案咨询

实现逻辑说明 #

你的需求本质是要找到所有和初始标记值存在关联的节点，属于连通集合遍历场景：

• 初始标记值：n3 == "a"对应行的n1值
• 关联规则：只要某行的n1或n2在已标记值中，该行的n1就加入标记集合，支持链式传递
• 最终所有n1在标记集合的行，n4统一填充为red

测试数据 #

df <- data.frame(n1 = c(1,2,1,2,5,6,8,9,8,8),
                 n2 = c(100,1000,500,1,NA,NA,2,8,10,15),
                 n3 = c("a", "a", "a", NA, "b", "c",NA,NA,NA,NA),
                 n4 = c("red", "red", NA, NA, NA, NA,NA,NA,NA,NA))

实现方案 #

1. tidyverse 版本

library(tidyverse)

# 提取初始标记种子值
red_values <- df %>% filter(n3 == "a") %>% pull(n1) %>% unique()
old_len <- 0

# 链式遍历新增关联值，直到无新值加入为止
while(old_len != length(red_values)){
  old_len <- length(red_values)
  new_vals <- df %>% 
    filter(n1 %in% red_values | n2 %in% red_values) %>% 
    pull(n1) %>% 
    unique()
  red_values <- unique(c(red_values, new_vals))
}

# 填充n4列
df <- df %>% mutate(n4 = ifelse(n1 %in% red_values, "red", n4))

2. base R 版本

# 提取初始标记种子值
red_values <- unique(df$n1[df$n3 == "a"])
old_len <- 0

# 链式遍历新增关联值，直到无新值加入为止
while(old_len != length(red_values)){
  old_len <- length(red_values)
  idx <- df$n1 %in% red_values | df$n2 %in% red_values
  new_vals <- unique(df$n1[idx])
  red_values <- unique(c(red_values, new_vals))
}

# 填充n4列
df$n4[df$n1 %in% red_values] <- "red"

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