Select N rows above and below match

Question

I would like to select N rows above and below a match.

I'm trying the command:

mtcars[which(mtcars$vs == 1) + c(-1:1), ]

It returns the follow warning:

Warning message: In which(mtcars$vs == 1) + c(-1:1): longer object length is not a multiple of shorter object length

Answer 1

We can write a short function to return all elements of vec that match val or that are within n elements (either direction):

newfun <- function(vec, val, n) {
    rows <- sapply(which(vec==val), function(x) seq(x-n, x+n, 1))
    rows <- unique(sort(rows[rows>0 & rows<length(vec)]))
    return(vec[rows])
}

For example:

newfun(mtcars$vs, 1, 2)

Answer 2

Before adding the desired range of indices to your focal index (the result from which), you need to repeat each value to the length of your range.

# set the number of values to select, before and after each focal index
n <- 1

# create a range of (relative) indices
i <- -n:n

# repeat focal indices 
# add range of n prior and following indices  
ix <- rep(which(mtcars$vs == 1), each = length(i)) + i

# select unique indices, truncated to the relevant range of rows,...
unique(ix[ix > 0 & ix <= nrow(mtcars)])
# [1]  2  3  4  5  6  7  8  9 10 11 12 17 18 19 20 21 22 25 26 27 28 29 31 32

# ...which then can be used to subset data
mtcars[unique(ix[ix > 0 & ix <= nrow(mtcars)]), ] 

Answer 3

This seems to be a simple question but is not as trivial as presumably expected.

The issue is that which(mtcars$vs == 1) returns a vector rather than a single value:

[1]  3  4  6  8  9 10 11 18 19 20 21 26 28 32

If another vector -1:1 (which is c(-1L, 0L, 1L)) is added to it, the normal R rules for operations on vectors of unequal lengths apply: The recycling rule says

Any short vector operands are extended by recycling their values until they match the size of any other operands.

Therefore the shorter vector -1:1 will be recycled to the length of which(mtcars$vs == 1), i.e.,

rep(-1:1, length.out = length(which(mtcars$vs == 1)))

 [1] -1  0  1 -1  0  1 -1  0  1 -1  0  1 -1  0

Therefore, the result of

which(mtcars$vs == 1) + -1:1

is the element-wise sum of the elements of both vectors where the shorter vector has been recycled to match the length of the longer vector.

 [1]  2  4  7  7  9 11 10 18 20 19 21 27 27 32

which is propably not what the OP has expected.

In addition, we get the

Warning message:
In which(mtcars$vs == 1) + -1:1 :
longer object length is not a multiple of shorter object length

because which(mtcars$vs == 1) has length 14 and -1:1 has length 3.

Solution using outer()

In order to select the N rows above and below each matching row, we need to add -N:N to each row number returned by which(mtcars$vs == 1):

outer(which(mtcars$vs == 1), -1:1, `+`)

      [,1] [,2] [,3]
 [1,]    2    3    4
 [2,]    3    4    5
 [3,]    5    6    7
 [4,]    7    8    9
 [5,]    8    9   10
 [6,]    9   10   11
 [7,]   10   11   12
 [8,]   17   18   19
 [9,]   18   19   20
[10,]   19   20   21
[11,]   20   21   22
[12,]   25   26   27
[13,]   27   28   29
[14,]   31   32   33

Now, we have an array of all row numbers. Unfortunately, it cannot be used directly for subsetting because it contains duplicates and there are row numbers which do not exist in mtcars. So the the result has to be "post-processed" before it can be used for subsetting.

library(magrittr) # piping used for clarity
rn <- outer(which(mtcars$vs == 1), -1:1, `+`) %>% 
  as.vector() %>% 
  unique() %>% 
  Filter(function(x) x[1 <= x & x <= nrow(mtcars)], .)

rn

 [1]  2  3  4  5  6  7  8  9 10 11 12 17 18 19 20 21 22 25 26 27 28 29 31 32

mtcars[rn, ]

                   mpg cyl  disp  hp drat    wt  qsec vs am gear carb
Mazda RX4 Wag     21.0   6 160.0 110 3.90 2.875 17.02  0  1    4    4
Datsun 710        22.8   4 108.0  93 3.85 2.320 18.61  1  1    4    1
Hornet 4 Drive    21.4   6 258.0 110 3.08 3.215 19.44  1  0    3    1
Hornet Sportabout 18.7   8 360.0 175 3.15 3.440 17.02  0  0    3    2
Valiant           18.1   6 225.0 105 2.76 3.460 20.22  1  0    3    1
Duster 360        14.3   8 360.0 245 3.21 3.570 15.84  0  0    3    4
Merc 240D         24.4   4 146.7  62 3.69 3.190 20.00  1  0    4    2
Merc 230          22.8   4 140.8  95 3.92 3.150 22.90  1  0    4    2
Merc 280          19.2   6 167.6 123 3.92 3.440 18.30  1  0    4    4
Merc 280C         17.8   6 167.6 123 3.92 3.440 18.90  1  0    4    4
Merc 450SE        16.4   8 275.8 180 3.07 4.070 17.40  0  0    3    3
Chrysler Imperial 14.7   8 440.0 230 3.23 5.345 17.42  0  0    3    4
Fiat 128          32.4   4  78.7  66 4.08 2.200 19.47  1  1    4    1
Honda Civic       30.4   4  75.7  52 4.93 1.615 18.52  1  1    4    2
Toyota Corolla    33.9   4  71.1  65 4.22 1.835 19.90  1  1    4    1
Toyota Corona     21.5   4 120.1  97 3.70 2.465 20.01  1  0    3    1
Dodge Challenger  15.5   8 318.0 150 2.76 3.520 16.87  0  0    3    2
Pontiac Firebird  19.2   8 400.0 175 3.08 3.845 17.05  0  0    3    2
Fiat X1-9         27.3   4  79.0  66 4.08 1.935 18.90  1  1    4    1
Porsche 914-2     26.0   4 120.3  91 4.43 2.140 16.70  0  1    5    2
Lotus Europa      30.4   4  95.1 113 3.77 1.513 16.90  1  1    5    2
Ford Pantera L    15.8   8 351.0 264 4.22 3.170 14.50  0  1    5    4
Maserati Bora     15.0   8 301.0 335 3.54 3.570 14.60  0  1    5    8
Volvo 142E        21.4   4 121.0 109 4.11 2.780 18.60  1  1    4    2