Efficiently checking value of other row in data.table

Question

Note: This is a question that I originally posted to the data.table help group. Matt Dowle asked for a more detailed example and I posted this one, but I had trouble with the formatting in e-mail. I already know how to format on SO, so I thought I would post it here instead. What I am basically trying to do is subset rows from a data.table based on a value in that row as well as a value in a preceding or following row. Right now, I am creating new columns for future and past rows and then keying the data.table on these columns, but this is resource-intensive and onerous. The below example illustrates the approach I am using now. The example uses words in documents (I use numeric indices for both). I want to subset for a particular word, but only if it is preceded or followed by another word or set of words: I first create a dummy dataset with ten documents containing one million words. There are three unique words in the set. ``` library(data.table) set.seed(1000) DT<-data.table(wordi…

Accepted Answer

Here's another idea that just sprung to my mind. It requires just creating one more column and uses binary search for subset.

On the `DT` you've generated from your data, first we'll add the extra column:

```
# the extra column:
DT[, I := .I]
```

We need this because we'll `setkey` on `docindex` and `wordindex`. This is the only way we can subset without creating extra columns (at least what I could think of). So, we'll need a way to extract the "original" positions later to check for your condition (hence the `I`).

After adding the extra column, let's set the key on the two columns mentioned above:

```
setkey(DT, docindex, wordindex)
```

Great! The idea from here is to extract the positions where your desired word matches - here that value is `1L`. Then, extract all the other words you may (or may not) want to come after this word at the right position. Then, we simply keep (or remove) those indices that satisfy the condition.

Here's a function that'll take care of this. It is by no means complete, but should give you an idea.

```
foo <- function(DT, doc_key, word_key, rest_key=NULL, match=FALSE) {
    ## note that I'm using 1.9.3, where this results in a vector
    ## if you're using 1.9.2, you'll have to change the joins accordingly
    idx1 = DT[J(doc_key, word_key), I]
    for (i in seq_along(rest_key)) {
        this_key = rest_key[i]
        idx2 = DT[J(doc_key, this_key), I]
        if (match) idx1 = idx1[which((idx1+i) %in% idx2)]
        else idx1 = idx1[which(!(idx1+i) %in% idx2)]
    }
    DT[idx1, .N, by=c(key(DT)[1L])]
}
```

Here, `DT` is the `data.table` to which `I` column has been added, and then `setkey` has been called on the two columns as mentioned before.

`doc_key` basically contains all the unique values in `docindex` - here 1:10. `word_key` is basically 1L here. `rest_key` is the values you'd like to check does not occur at `i`th position after the position of `word_key`.

First we extract `I` for all matches of `1L` in `idx1` (straightforward). Next, we loop through each value of `rest_key` and add that position to `idx1` = `idx1+i` and check if that value occurs in `idx2`. If so, based of whether you like to extract matching or non-matching entries, we'll keep (or remove them).

And at the end of this loop, `idx1` should have only the desired entries. Hope this helps. Shown below is a demonstration of the cases already discussed in the other answer.

Let's consider your first scenario. Count of all entries, for each group in `docindex` where ith position is `1L` and `i+1`th is not 2L. This is basically:

```
system.time(ans1 <- foo(DT, 1:10, 1L, 2L, FALSE))

#  user  system elapsed 
# 0.066   0.019   0.085

# old method took 0.12 seconds

#     docindex     N
#  1:        1 22301
#  2:        2 21836
#  3:        3 22491
#  4:        4 21831
#  5:        5 22218
#  6:        6 21914
#  7:        7 22370
#  8:        8 22265
#  9:        9 22211
# 10:       10 22190
```

What about the second scenario? Here, we'd like the the `i+1`th and `i+2`th position to be 2L and 3L, as opposed to the not equal scenario in the earlier case. So, we set `match=TRUE` here.

```
system.time(ans2 <- foo(DT, 1:10, 1L, 2:3,TRUE))
#  user  system elapsed 
# 0.080   0.011   0.090

# old method took 0.22 seconds

#     docindex    N
#  1:        1 3684
#  2:        2 3746
#  3:        3 3717
#  4:        4 3727
#  5:        5 3700
#  6:        6 3779
#  7:        7 3702
#  8:        8 3756
#  9:        9 3702
# 10:       10 3744
```

It's easy to expand this function. For ex: if you'd like to have `i+1`th to be equal to `2L` but `i+2`th not equal to `3L` then, you can change `match` to be a vector = `length(rest_key)` specifying corresponding logical values.

I hope this is fast for your actual case - at least faster than the other version.

HTH

Efficiently checking value of other row in data.table

Solution

Problem