5

Given even n, I would like to make a 3^(n/2)-1 by n 2d-numpy array. The first half of every row should iterate over all possible values of -1,0,1 and the second half should be zero. However, the first half should never be all zeros.

This code almost works except it includes the all zero row which I don't want.

n = 4
M = [list(row) +[0]*(n/2) for row in itertools.product([-1,0,1], repeat = n/2)]
print np.array(M)

It gives

[[-1 -1  0  0]
 [-1  0  0  0]
 [-1  1  0  0]
 [ 0 -1  0  0]
 [ 0  0  0  0]
 [ 0  1  0  0]
 [ 1 -1  0  0]
 [ 1  0  0  0]
 [ 1  1  0  0]]

Is there a less horrible and more time and space-efficient way to do this? n will ultimately be 30 and I won't be printing it out of course. 3^15 is only 14,348,907 and yet the code uses all the RAM on my 8GB machine when I set n=30 as well as taking a very long time.

How can I make the numpy array directly without going via itertools, lists etc.?

user92234
  • 106
  • 2
  • 10
Simd
  • 19,447
  • 42
  • 136
  • 271

1 Answers1

5

This will create your array without needing any large auxiliary memory allocation:

n = 30
assert 2 * (n // 2) == n
rows = 3**(n//2)
cols = n

arr = np.zeros((rows, cols), dtype=int)
shape = (rows,)

source = np.array([-1, 0, 1], dtype=np.int)[:, None]

for col in range(n//2):
    shape = (-1, 3, shape[-1]//3,)
    col_view = arr[:, col]
    col_view.shape = shape
    col_view[:] = source

Takes about 10 sec to complete on my laptop. It is loosely based on this great answer to a similar question.

What can't be done easily with this approach is getting rid of the middle line of zeros during construction, as it is central to the workings of the algorithm. You can always get rid of it afterwards:

arr = np.delete(arr, rows//2, axis=0)

But this will allocate a new array and copy the contents to it before discarding the old one, so your memory requirements will all of a sudden double.

Can't think of any easy, fast way around it with Python or NumPy.

Community
  • 1
  • 1
Jaime
  • 65,696
  • 17
  • 124
  • 159
  • I did an in-place version using the Cartesian-product-as-base-expansion trick, dividing into two steps to handle the zero, but even after adding `np.delete` to yours your code is *still* ~2x faster. – DSM Jul 16 '15 at 16:38
  • In this case it really helps to choose a smaller dtype and use cache friendly ordering. Because you completely got rid of the Python overhead, the only thing that remains is raw memory access speed. –  Jul 16 '15 at 21:06
  • Indeed, making the result array be Fortran order will probably have a big positive impact on performance. – Jaime Jul 16 '15 at 21:14