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

@AmirhosseinKiani you mean changing range(1,l+1) to range(1,n,6)? — Rararat
– Rararat, Commented Mar 23, 2022 at 22:37
Kinda, you probably need to generate two ranges and eliminate the first element of each. range is not comparable to the numpy answers. I suggest you to use them, and do not consider range at all:) — TheFaultInOurStars
– TheFaultInOurStars, Commented Mar 23, 2022 at 22:41

user7864386 · Accepted Answer · 2022-03-23 22:34:25Z

3

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

answered Mar 23, 2022 at 22:34

user7864386

Sign up to request clarification or add additional context in comments.

Comments

Tim Roberts · Accepted Answer · 2022-03-23 22:30:19Z

3

Certainly.

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

answered Mar 23, 2022 at 22:30

Tim Roberts

55.3k4 gold badges29 silver badges41 bronze badges

1 Comment

Rararat Over a year ago

Thanks! But, would there be a way to make this in order without recurring to sorting?

score 2 · Accepted Answer · 2022-03-23 23:58:26Z

2

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!

edited Mar 23, 2022 at 23:58

answered Mar 23, 2022 at 22:56

user17242583

Collectives™ on Stack Overflow

Python Array Computations: get two values per item

3 Answers 3

Comments

1 Comment

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

3 Answers 3

Comments

1 Comment

Comments

Your Answer

Sign up or log in

Post as a guest

Related