是否有一种 pythonic 方法从列表或 numpy 数组中采样 N 个连续元素

Question

是否有 pythonic 方法从列表或 numpy 数组中 select N 个连续元素。

所以假设：

Choice = [1,2,3,4,5,6]

我想创建一个长度为 N 的新列表，方法是随机 selectChoice 中的元素 X 以及选择后的 N-1 个连续元素。

所以如果：

X = 4 
N = 4

结果列表为：

Selection = [5,6,1,2]

我认为类似下面的东西会起作用。

S = [] 
for i in range(X,X+N):
    S.append(Selection[i%6])

但我想知道是否有 python 或 numpy 函数可以一次 select 更高效的元素。

Answer 1

您可以使用列表理解，对索引使用模运算以使其保持在列表范围内：

Choice = [1,2,3,4,5,6] 
X = 4 
N = 4
L = len(Choice)
Selection = [Choice[i % L] for i in range(X, X+N)]
print(Selection)

输出

[5, 6, 1, 2]

注意如果N小于等于len(Choice)，可以大大简化代码：

Choice = [1,2,3,4,5,6] 
X = 4 
N = 4
L = len(Choice)
Selection = Choice[X:X+N] if X+N <= L else Choice[X:] + Choice[:X+N-L]
print(Selection)

Answer 2

使用 itertools，特别是 islice 和 cycle。

start = random.randint(0, len(Choice) - 1)
list(islice(cycle(Choice), start, start + n))

cycle(Choice) 是一个重复原始列表的无限序列，因此切片 start:start + n 将在必要时换行。

Answer 3

这是一个 numpy 方法：

import numpy as np

Selection = np.take(Choice, range(X,N+X), mode='wrap')

即使 Choice 是一个 Python 列表而不是 numpy 数组也能工作。

Answer 4

由于您要求的是最有效的方法，我创建了一个小基准来测试该线程中提出的解决方案。

我将您当前的解决方案重写为：

def op(choice, x):
    n = len(choice)
    selection = []
    for i in range(x, x + n):
        selection.append(choice[i % n])
    return selection

其中 choice 是输入列表，x 是随机索引。

这些是如果 choice 包含 1_000_000 个随机数的结果：

chepner: 0.10840400000000017 s
nick: 0.2066781999999998 s
op: 0.25887470000000024 s
fountainhead: 0.3679908000000003 s

完整代码

import random
from itertools import cycle, islice
from time import perf_counter as pc
import numpy as np


def op(choice, x):
    n = len(choice)
    selection = []
    for i in range(x, x + n):
        selection.append(choice[i % n])
    return selection


def nick(choice, x):
    n = len(choice)
    return [choice[i % n] for i in range(x, x + n)]


def fountainhead(choice, x):
    n = len(choice)
    return np.take(choice, range(x, x + n), mode='wrap')


def chepner(choice, x):
    n = len(choice)
    return list(islice(cycle(choice), x, x + n))


results = []
n = 1_000_000
choice = random.sample(range(n), n)
x = random.randint(0, n - 1)

# Correctness
assert op(choice, x) == nick(choice,x) == chepner(choice,x) == list(fountainhead(choice,x))

# Benchmark
for f in op, nick, chepner, fountainhead:
    t0 = pc()
    f(choice, x)
    t1 = pc()
    results.append((t1 - t0, f))

for t, f in sorted(results):
    print(f'{f.__name__}: {t} s')

Answer 5

如果使用 numpy 数组作为源，我们当然可以使用 numpy “花式索引”。

因此，如果 ChoiceArray 是列表 Choice 的 numpy 数组等效项，并且如果 L 是 len(Choice) 或 len(ChoiceArray):

Selection = ChoiceArray [np.arange(X, N+X) % L]

是否有一种 pythonic 方法从列表或 numpy 数组中采样 N 个连续元素

Is there a pythonic way to sample N consecutive elements from a list or numpy array

python

numpy

sample

完整代码