基于data.table高效判断数据框两列间的传递性配对问题

原始数据 #

df <- structure(list(subject = c(1003L, 1003L, 1003L, 1003L, 1003L, 
1003L, 1003L, 1003L, 1003L, 1003L, 1003L), stim1 = c(50L, 48L, 
49L, 47L, 47L, 46L, 50L, 48L, 46L, 50L, 54L), stim2 = c(51L, 
50L, 51L, 49L, 46L, 48L, 48L, 51L, 50L, 49L, 59L), feedback = c(1L, 
1L, 1L, 1L, 1L, 1L, 1L, 0L, 0L, 0L, 0L)), row.names = c(1L, 
2L, 3L, 4L, 5L, 6L, 10L, 428L, 433L, 434L, 435L), class = "data.frame")

需求说明 #

需要新增一列transitive_pair，规则如下：

• 按subject列分组处理
• 仅对feedback == 0的行计算该列值，其余行设为NaN
• 判断该行的stim1与stim2是否能通过feedback == 1的配对形成传递链，返回布尔值

示例说明 #

row 428: stim1=48 和 stim2=51
48和51未直接配对，但51与50配对（行1）、50与48配对（行10），因此transitive_pair[428] = TRUE

row 433: stim1=46 和 stim2=50
46与48配对（行6）、48与50配对（行2），因此transitive_pair[433] = TRUE

row 435: stim1=54 和 stim2=59
没有配对链能连接二者（59未在feedback==1的行中出现），因此transitive_pair[435] = FALSE

期望输出 #

structure(list(subject = c(1003L, 1003L, 1003L, 1003L, 1003L, 
1003L, 1003L, 1003L, 1003L, 1003L, 1003L), stim1 = c(50L, 48L, 
49L, 47L, 47L, 46L, 50L, 48L, 46L, 50L, 54L), stim2 = c(51L, 
50L, 51L, 49L, 46L, 48L, 48L, 51L, 50L, 49L, 59L), feedback = c(1L, 
1L, 1L, 1L, 1L, 1L, 1L, 0L, 0L, 0L, 0L), transitive_pair = c(NaN, 
NaN, NaN, NaN, NaN, NaN, NaN, 1, 1, 1, 0)), row.names = c(1L, 
2L, 3L, 4L, 5L, 6L, 10L, 428L, 433L, 434L, 435L), class = "data.frame")

这个问题本质是判断图中两个节点是否连通：把feedback==1的每对stim1和stim2看作无向边，每个stim值是节点，传递链存在等价于两个节点在同一个连通分量中。可以用igraph包实现：

步骤1：安装并加载所需包 #

install.packages("igraph")
library(igraph)
library(dplyr)

步骤2：分组处理并计算transitive_pair #

df_result <- df %>%
  group_by(subject) %>%
  mutate(
    transitive_pair = case_when(
      feedback != 0 ~ NaN,
      TRUE ~ {
        # 提取当前subject下feedback==1的配对，构建无向图
        edges <- filter(cur_data(), feedback == 1) %>% select(stim1, stim2)
        graph <- graph_from_data_frame(edges, directed = FALSE)
        # 检查stim1和stim2是否在同一连通分量
        map2_lgl(stim1, stim2, function(s1, s2) {
          # 如果任一节点不在图中，直接返回FALSE
          if (!s1 %in% V(graph)$name || !s2 %in% V(graph)$name) {
            return(FALSE)
          }
          # 检查连通性
          components(graph)$membership[as.character(s1)] == components(graph)$membership[as.character(s2)]
        })
      }
    )
  ) %>%
  ungroup()

验证结果 #

运行上述代码后，df_result的transitive_pair列将与期望输出一致。

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