Extract reachable rows based on id

user_7120667 0
Viewed 58

I am trying to write a function which will return all rows related to that id.

Input Data: 

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')
      )

E.g. 

sample_rows <- function(id1){

  #code here 
  return(output)
}

Desired Output 1:

sample_rows(id1='A1') should return following data frame:

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

Desired Output 2:

sample_rows(id1 = 'C3') should return following data frame:

  id1    id2
  C3     KK3
  B3     KK3
data.table dplyr tidyverse plyr r
2 Answers

Don't use sample as function name, you risk overriding R own function

base R solution: 

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 solution:

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

Created on 2020-04-03 by the reprex package (v0.3.0)

user_12612810

One solution utilizing igraph and dplyr could be:

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

Or:

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])
}
user_59645570
Related