Выбрать ближайшие x элементов по индексу в списке / векторе

num_closest_by_indices <- function(v, idx, num) {
  # Try the base case, where idx is not within (num/2) of the edge
  i <- abs(seq_along(x) - idx)
  i[idx] <- +Inf # sentinel

  # If there are not enough elements in the base case, incrementally add more
  for (cutoff_idx in seq(floor(num/2), num)) {
    if (sum(i <= cutoff_idx) >= num) {
      # This will add two extra indices every iteration. Strictly if we have an even length, we should add the leftmost one first and `continue`, to break ties towards the left.
      return(v[i <= cutoff_idx])
    }
  }
} 

Вот иллюстрация этого алгоритма: мы ранжируем индексы в порядке желательности, а затем выбираем самые низкие из легальных для num:

> seq_along(x)
  1 2 3 4 5 6 7 8 9
> seq_along(x) - idx
  -2 -1  0  1  2  3  4  5  6
> i <- abs(seq_along(x) - idx)
   2  1  0  1  2  3  4  5  6
> i[idx] <- +Inf # sentinel to prevent us returning the element itself
   2   1 Inf   1   2   3   4   5   6

Теперь мы можем просто найти элементы num с наименьшими значениями (разорвать связи произвольно, если у вас нет предпочтений (слева)). Наше первое предположение - все индексы <= (num / 2); этого может быть недостаточно, если index находится в пределах (num/2) начала / конца.

> i <= 2
  TRUE  TRUE FALSE  TRUE  TRUE FALSE FALSE FALSE FALSE
> v[i <= 2]
  1 2 4 5

Итак, адаптация кода @ dash2 для обработки угловых случаев, когда некоторые индексы недопустимы (неположительные или> length (x)), то есть ! %in% 1:L. Тогда min(elems) будет числом нелегальных индексов, которые мы не можем выбрать, следовательно, мы должны выбрать больше abs(min(elems)).

Заметки:

• в итоге код проще и быстрее обрабатывать его по трем частям. Ой.
• на самом деле кажется, что это упрощает ситуацию, если мы выбираем индексы (num+1), а затем удаляем idx перед возвратом ответа. Использование result[-idx] для его удаления.

Вот так:

f <- function (vec, elem, n) {
  elems <- seq(elem - ceiling(n/2), elem + floor(n/2))
  if (max(elems) > length(vec)) elems <- elems - (max(elems) - length(vec))
  if (elems[1] < 1) elems <- elems + (1 - elems[1])
  elems <- setdiff(elems, elem)
  vec[elems]
}

Давать результаты:

> f(1:9, 1, 5)
[1] 2 3 4 5 6
> f(1:9, 9, 5)
[1] 4 5 6 7 8
> f(1:9, 2, 5)
[1] 1 3 4 5 6
> f(1:9, 4, 5)
[1] 1 2 3 5 6
> f(1:9, 4, 4)
[1] 2 3 5 6
> f(1:9, 2, 4)
[1] 1 3 4 5
> f(1:9, 1, 4)
[1] 2 3 4 5
> f(1:9, 9, 4)
[1] 5 6 7 8

Сначала запустите функцию с переменным аргументом x, а после - ссылками table и n.

.nearest_n <- function(x, table, n) {

Алгоритм предполагает, что table является числовым, без каких-либо дубликатов и все значения конечны; n должен быть меньше или равен длине стола.

    ## assert & setup
    stopifnot(
        is.numeric(table), !anyDuplicated(table), all(is.finite(table)),
        n <= length(table)
    )

Отсортируйте таблицу, а затем «зафиксируйте» максимальные и минимальные значения.

    ## sort and clamp
    table <- c(-Inf, sort(table), Inf)
    len <- length(table)

Найдите интервал в table, где встречается x; findInterval() использует эффективный поиск. Используйте индекс интервала в качестве начального нижнего индекса и добавьте 1 для верхнего индекса, следя за тем, чтобы он оставался в границах.

    ## where to start?
    lower <- findInterval(x, table)
    upper <- min(lower + 1L, len)

Найдите ближайших соседей n, сравнив расстояние между нижним и верхним индексами и x, запишите ближайшее значение и увеличьте нижний или верхний индекс в зависимости от ситуации, следя за тем, чтобы они оставались в пределах

    ## find
    nearest <- numeric(n)
    for (i in seq_len(n)) {
        if (abs(x - table[lower]) < abs(x - table[upper])) {
            nearest[i] = table[lower]
            lower = max(1L, lower - 1L)
        } else {
            nearest[i] = table[upper]
            upper = min(len, upper + 1L)
        }
    }

Затем верните решение и завершите функцию

    nearest
}

Код может показаться многословным, но на самом деле он относительно эффективен, поскольку единственные операции со всем вектором (sort(), findInterval()) эффективно реализованы в R.

Особым преимуществом этого подхода является то, что его можно векторизовать в своем первом аргументе, вычисляя тест для использования нижнего (use_lower = ...) в качестве вектора и использования pmin() / pmax() в качестве зажимов.

.nearest_n <- function(x, table, n) {
    ## assert & setup
    stopifnot(
        is.numeric(table), !anyDuplicated(table), all(is.finite(table)),
        n <= length(table)
    )

    ## sort and clamp
    table <- c(-Inf, sort(table), Inf)
    len <- length(table)

    ## where to start?
    lower <- findInterval(x, table)
    upper <- pmin(lower + 1L, len)

    ## find
    nearest <- matrix(0, nrow = length(x), ncol = n)
    for (i in seq_len(n)) {
        use_lower <- abs(x - table[lower]) < abs(x - table[upper])
        nearest[,i] <- ifelse(use_lower, table[lower], table[upper])
        lower[use_lower] <- pmax(1L, lower[use_lower] - 1L)
        upper[!use_lower] <- pmin(len, upper[!use_lower] + 1L)
    }

    # return
    nearest
}

Например

> set.seed(123)
> table <- sample(100, 10)
> sort(table)
 [1]  5 29 41 42 50 51 79 83 86 91
> .nearest_n(c(30, 20), table, 4)
     [,1] [,2] [,3] [,4]
[1,]   29   41   42   50
[2,]   29    5   41   42

Обобщите это, взяв любой аргумент и приведя его к требуемой форме с помощью справочной таблицы table0 и индексов в ней table1.

nearest_n <- function(x, table, n) {
    ## coerce to common form
    table0 <- sort(unique(c(x, table)))
    x <- match(x, table0)
    table1 <- match(table, table0)

    ## find nearest
    m <- .nearest_n(x, table1, n)

    ## result in original form
    matrix(table0[m], nrow = nrow(m))
}

В качестве примера...

> set.seed(123)
> table <- sample(c(letters, LETTERS), 30)
> nearest_n(c("M", "Z"), table, 5)
     [,1] [,2] [,3] [,4] [,5]
[1,] "o"  "L"  "O"  "l"  "P" 
[2,] "Z"  "z"  "Y"  "y"  "w"