Identify neighbors in a data.table

user_8099679 0
Viewed 77

I have the following data.table*:

dt <- data.table(
  LEFT  = letters[c(1:5, 10:12)],
  SELF  = letters[c(2:6, 11:13)],
  RIGHT = letters[c(3:7, 12:14)]
  )

dt <- dt[sample(nrow(dt)), ]
dt[]

#    LEFT SELF RIGHT
# 1:    e    f     g
# 2:    d    e     f
# 3:    j    k     l
# 4:    c    d     e
# 5:    l    m     n
# 6:    b    c     d
# 7:    a    b     c
# 8:    k    l     m

LEFT and RIGHT indicate the neighbors of SELF. I'd like to index the table such that I group and order contiguous neighbors. An ideal output could look something like this:

dt_out[]
#    LEFT SELF RIGHT RUN ORDER
# 1:    e    f     g   1     5
# 2:    d    e     f   1     4
# 3:    j    k     l   2     1
# 4:    c    d     e   1     3
# 5:    l    m     n   2     3
# 6:    b    c     d   1     2
# 7:    a    b     c   1     1
# 8:    k    l     m   2     2

run_1 <- dt_out[RUN == 1][order(ORDER)][["SELF"]]
run_1
# [1] "b" "c" "d" "e" "f"

I am tempted to write a function to apply to SELF to identify where LEFT == SELF == RIGHT, but I think this is the wrong avenue to go down given data.table's order-by-group capabilities.

*in reality, my data.table has 1.9M observations and is not ordered in any meaningful way.

data.table r
2 Answers

We can use rleid together with %/% 2 to create the RUN variable and then we can just group by RUN and get the row number with 1:.N to create the ORDER column.

library(data.table)

dt[,
   RUN := rleid(shift(RIGHT, type = "lag", fill = SELF[1]) == SELF) %/% 2 + 1
][,
  ORDER := 1:.N,
  by = RUN][]

#>      LEFT   SELF  RIGHT   RUN ORDER
#>    <char> <char> <char> <num> <int>
#> 1:      a      b      c     1     1
#> 2:      b      c      d     1     2
#> 3:      c      d      e     1     3
#> 4:      d      e      f     1     4
#> 5:      e      f      g     1     5
#> 6:      j      k      l     2     1
#> 7:      k      l      m     2     2
#> 8:      l      m      n     2     3

Created on 2022-03-30 by the reprex package (v0.3.0)

user_93493020

Here is some playing around with recursion. If your connected components tend to be big this might fail (recursion limit) but maybe you could adapt the logic.

dt[, grp   := NA_character_]
dt[, order := NA_integer_]
setkey(dt, SELF)

follow_path = function(node, order = 0L, grp = NULL) {
  i = dt[list(node), which = TRUE]
  if (is.na(i) || !is.na(dt$grp[i])) return();
  if (is.null(grp)) {
    set(dt, i, "grp", dt$SELF[i])
    set(dt, i, "order", 0L)
  } else {
    set(dt, i, "grp", grp)
    set(dt, i, "order", order)    
  }
  follow_path(dt$RIGHT[i], order = order +1, grp = dt$grp[i])
  follow_path(dt$LEFT[i],  order = order -1, grp = dt$grp[i])
}

for (i in 1L:nrow(dt)) follow_path(dt$SELF[i])
dt[, grp := as.integer(factor(grp))]
dt[, order := frank(order), by = grp]
dt

#      LEFT   SELF  RIGHT   grp order
#    <char> <char> <char> <int> <int>
# 1:      a      b      c     1     1
# 2:      b      c      d     1     2
# 3:      c      d      e     1     3
# 4:      d      e      f     1     4
# 5:      e      f      g     1     5
# 6:      j      k      l     2     1
# 7:      k      l      m     2     2
# 8:      l      m      n     2     3
user_45522950
Related