Извлечение достижимых строк на основе идентификатора

Question

Я пытаюсь написать function, который будет возвращать все строки, связанные с этим идентификатором.

Входные данные:

df <- data.frame(
         id1=c('A1','A1','B1','B2','C2','C3','B3','D1','D2','E1'),
         id2=c('P1','P2','P1','P2','P1','KK3','KK3','DL5','FD7','LO6')
      )

Например,

sample_rows <- function(id1){

  #code here 
  return(output)
}

Желаемый выход 1:

sample_rows(id1='A1') должен возвращать следующий кадр данных:

  id1    id2
  A1     P1
  A1     P2
  B1     P1
  B2     P2
  C2     P1

Желаемый выход 2:

sample_rows(id1 = 'C3') должен вернуть следующий фрейм данных:

  id1    id2
  C3     KK3
  B3     KK3

Answer 1

Не используйте sample в качестве имени функции, вы рискуете переопределить собственную функцию R

base R решение:

my_fun <- function(id1) {
  id2s <- df[df$id1 == id1, 'id2']
  df[df$id2 %in% id2s,]
}

my_fun(id1 = 'A1')
#>   id1 id2
#> 1  A1  P1
#> 2  A1  P2
#> 3  B1  P1
#> 4  B2  P2
#> 5  C2  P1

my_fun(id1 = 'C3')
#>   id1 id2
#> 6  C3 KK3
#> 7  B3 KK3

dplyr решение:

library(dplyr)

my_dplyr_fun <- function(data, id1) {
  id2s <- filter(data, id1 == {{id1}}) %>% 
    pull(id2)
  data %>%
    filter(id2 %in% id2s)
}

df %>% 
  my_dplyr_fun('C3')
#>   id1 id2
#> 1  C3 KK3
#> 2  B3 KK3

^{Создано в 2020-04-03 пакетом Представления (v0.3.0)}

Answer 2

Одно решение с использованием igraph и dplyr может быть:

fun <- function(data, id1) {
 data %>%
  left_join(stack(clusters(graph.data.frame(data))$membership),
           by = c("id2" = "ind")) %>%
 filter(values == values[id1 == !!id1])
}

fun(df, "A1")

  id1 id2 common_id
1  A1  P1         1
2  A1  P2         1
3  B1  P1         1
4  B2  P2         1
5  C2  P1         1

Или:

fun <- function(data, id1) {
 clust <- clusters(graph.data.frame(data))$membership
 data %>%
  mutate(common_id = clust[match(id2, names(clust))]) %>%
  filter(common_id == common_id[id1 == !!id1])
}