How to match observations that are within +/- 5 of each other in R?

user_12532227 0
Viewed 113

Say I have a dataframe that looks like the following:

dat <- data.frame("firstName" = c("John", "John", "Mary", "Bob", "Mary", "Bob"), "age"= c(21, 24, 35, 30, 20, 27))

I want to create a third variable dat$id that assign the same number if an observation's age is within +/- 5 years of another observation and has the same firstName. So the dataframe will look like this: 

dat <- data.frame("firstName" = c("John", "John", "Mary", "Bob", "Mary", "Bob"), "age"= c(21, 24, 35, 30, 20, 27), "id"= c(1,1,2,3,4,3))

I have a very large dataset of names and ages and would like to find a more automated way of assigning id's. I considered creating age bins for every 5 years from age 20, but this will not match observations who are in different bins, but still within 5 years of age.

string-matching r
3 Answers

1) sqldf/igraph Match each row to those rows having the same name, an age within 5 and the row is not itself. If there is no such match then match the row to itself so that all rows are accounted for. The rows and their matches can then be turned into an edgelist and subsequently an igraph, g. Find the connected components of and assign the membership ids to the rows of the original data frame.

In the example data each connected component is size 1 or 2 but this approach can handle any size and not just those.

library(igraph)
library(sqldf)

s <- sqldf("select a.rowid, a.*, b.rowid as match 
  from dat a left join dat b
    on a.firstname = b.firstname and 
      abs(a.age - b.age) < 5 and
      a.rowid != b.rowid")
e <- cbind(s$rowid, s$match) # edgelist
e[is.na(s$match), 2] <- e[is.na(s$match), 1]  
g <- graph_from_edgelist(e)
transform(dat, id = components(g)$membership)

giving:

  firstName age id
1      John  21  1
2      John  24  1
3      Mary  35  2
4       Bob  30  3
5      Mary  20  4
6       Bob  27  3

We can visualize the graph like this:

plot(g)

(continued after graph)

screenshot

2) Base R This solution is motivated, in part, by the other solutions but has significant advantages in that it only uses base R, only 2 lines of code, like (1) also handles connected components of any size, produces the correct answer and is fully vectorized. It works by sorting the data and then pulling forward the id or generating a new one depending on the condition shown. 

o <- with(dat, order(firstName, age))
transform(dat[o,], id = cumsum(c(1, diff(xtfrm(firstName)) | diff(age) > 5)))

giving:

  firstName age id
6       Bob  27  1
4       Bob  30  1
1      John  21  2
2      John  24  2
5      Mary  20  3
3      Mary  35  4
user_5165480

Without additional packages

dat <- data.frame("firstName" = c("John", "John", "Mary", "Bob", "Mary", "Bob"), "age"= c(21, 24, 35, 30, 20, 27))
n <- length(dat$firstName)

vals <- list()
for (i in 1:n) {
    fname <- dat$firstName[i]
    age <- dat$age[i]
    index <- which(fname == dat$firstName &
     (age > dat$age - 5) &
     (age < dat$age + 5))
    vals[[i]] <- index
}

vals <- unique(vals)
dat$id <- NA

for (i in 1:length(vals)) {
    dat$id[vals[[i]]] <- i
}

Result

  firstName age id
1      John  21  1
2      John  24  1
3      Mary  35  2
4       Bob  30  3
5      Mary  20  4
6       Bob  27  3
user_16239110

Here's an approach with lag from dplyr:

library(dplyr)
dat %>%
  group_by(firstName) %>%
  arrange(firstName,age) %>%
  mutate(id = cumsum(!(age - (lag(age,default = -Inf) ) <= 5)))
# A tibble: 6 x 3
# Groups:   firstName [3]
  firstName   age    id
  <fct>     <dbl> <int>
1 Bob          27     1
2 Bob          30     1
3 John         21     1
4 John         24     1
5 Mary         20     1
6 Mary         35     2
user_130953260
Related