Python Array Computations: get two values per item-CodePudding

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

CodePudding user response：

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)]

CodePudding user response：

Certainly.

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

CodePudding user response：

You can use np.dstack with two np.arange()s and a .ravel():

primeCandidates = np.dstack([np.arange(0, n-6, 6) 5, np.arange(0, n-6, 6) 7]).ravel()

Test:

>>> np.all(np.dstack([np.arange(0, n-6, 6) 5, np.arange(0, n-6, 6) 7]).ravel() == [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)