Is there a pythonic way to sample N consecutive elements from a list or numpy array
Use itertools
, specifically islice
and cycle
.
start = random.randint(0, len(Choice) - 1)
list(islice(cycle(Choice), start, start + n))
cycle(Choice)
is an infinite sequence that repeats your original list, so that the slice start:start + n
will wrap if necessary.
You could use a list comprehension, using modulo operations on the index to keep it in range of the list:
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)
Output
[5, 6, 1, 2]
Note that if N
is less than or equal to len(Choice)
, you can greatly simplify the code:
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)
Since you are asking for the most efficient way I created a little benchmark to test the solutions proposed in this thread.
I rewrote your current solution as:
def op(choice, x):
n = len(choice)
selection = []
for i in range(x, x + n):
selection.append(choice[i % n])
return selection
Where choice
is the input list and x
is the random index.
These are the results if choice
contains 1_000_000 random numbers:
chepner: 0.10840400000000017 s
nick: 0.2066781999999998 s
op: 0.25887470000000024 s
fountainhead: 0.3679908000000003 s
Full code
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')