Getting the square of the difference between many points in a NumPy array

Question

I have an array with 40000 numbers that are floats or ints. I need to perform some calculation.

To do this I have used nested for loop, but the code is really slow. Can I use something instead of nested the for loop? Is there any other way to reduce the execution time?
When I changed to a list comprehension the execution time reduce slightly. Why did the list comprehension take less time?

import numpy as np
import time as t

pox1=  np.random.randint(1000, size= 40000)
time = np.arange(40000)

y=np.zeros(len(pox1))
w=np.zeros(len(pox1))

start = t.time()
for num in range (len(time)-1): 
    z= [((pox1[i] - pox1[i-num]) ** 2) for i in range(num, len(pox1))]
    k=np.mean(z)
    y[num]=k

  # for i in range(num, len(pox1)):
  #    z.append((pox1[i] - pox1[i-num]) ** 2)


endtime = (t.time()-start)
  
print(y,endtime)

Answer 1

Python slow, C fast

By having the loops in Python you're basically cutting one of your legs off. It should come as no surprise that the one legged human will lose to the two legged human.

The question we should be asking is how can we slice pox1 to get what we want.

1. We should split the for loop into two comprehensions so we can see how we'd do it in numpy.

   a = [pox1[i]       for i in range(num, len(pox1))]
   b = [pox1[i - num] for i in range(num, len(pox1))]
   z = (a - b) ** 2
   

2. a is very simple we just slice with the values of the range. But b is a bit harder, we can just move the - num into the range and we'll be able to do the same.

   a = [pox1[i] for i in range(num, len(pox1))]
   b = [pox1[i] for i in range(0, len(pox1) - num)]
   z = (a - b) ** 2
   

3. Convert it to slices

   a = pox1[num : len(pox1)]
   b = pox1[0 : len(pox1) - num]
   z = (a - b) ** 2
   

4. Use some slice sugar

   a = pox1[num :]
   b = pox1[: len(pox1) - num]
   z = (a - b) ** 2
   

5. Put back in your code

   for num in range(len(time) - 1):
       y[num] = np.mean((pox1[num :] - pox1[: len(pox1) - num]) ** 2)
   

I changed the size to 4000 and this cut it from 5.04s to 0.07s. To get 40000 changes with the updated code takes 3.30s.

import numpy as np
import time as t

SIZE = 4000
pox1 = np.random.randint(1000, size=SIZE)
time = np.arange(SIZE)

y=np.zeros(len(pox1))
w=np.zeros(len(pox1))

start = t.time()

for num in range(len(time)-1):
    y[num] = np.mean((pox1[num :] - pox1[:len(pox1) - num]) ** 2)

endtime = t.time() - start
print(y, endtime)