Python Array Computations: get two values per item

Question

I'm trying to get a list with all numbers that are in the form 6n+1 or 6n-1. Currently I have this:

n = 100000000
l = int(n/6)
f1 = lambda x: (6*x)-1
f3 = lambda x: (6*x)+1
primeCandidate = [f(i) for i in range(1,l+1) for f in (f1,f3)]

This works nicely, and it gets me 2 values on the list per i, but I was wondering if I could do something similar with NumPy arrays

Answer 1

How about this. When you divide a number by 6, the modulo has to be either 1 or 5:

arr = np.arange(2, n)
out = arr[np.isin(np.mod(arr, 6), [1,5])]

Test:

assert arr[np.isin(np.mod(arr, 6), [1,5])].tolist() == [f(i) for i in range(1,l+1) for f in (f1,f3)]

Answer 2

Certainly.

pc1 = np.arange(0,n,6)+5
pc2 = np.arange(0,n,6)+1
pc = np.concatenate((pc1,pc2))

Answer 3

You can use np.dstack with two np.arange()s and a .reshape(-1):

primeCandidates = np.dstack((np.arange(5,n,6), np.arange(7,n,6))).reshape(-1)

Test:

>>> np.all(np.dstack((np.arange(5,n,6), np.arange(7,n,6))).reshape(-1) == [f(i) for i in range(1,l+1) for f in (f1,f3)])
True

This seems to beat enke's in terms of performance by about 10 times:

%timeit enke()
3.22 s ± 116 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit richardec()
259 ms ± 57.7 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

Credit goes to Michael Szczesny for suggesting to use reshape(-1) instead of ravel for speed, and for cleaning up the starts, stops, and steps of the ranges!