CODEWITHSUNDEEP
pythonpython-3.xnumpy
Waffle
Waffle

Reputation: 43

Add 2-d array to 3-d array with constantly changing index fast

I'm trying to add a 2-d array to a 3-d array with constantly changing index , I come up with following code:

import numpy as np

a = np.zeros([8, 3, 5])
k = 0
for i in range(2):
    for j in range(4):
        a[k, i: i + 2, j: j + 2] += np.ones([2, 2], dtype=int)
        k += 1
print(a)

which will give exactly what i want:

[[[1. 1. 0. 0. 0.]
  [1. 1. 0. 0. 0.]
  [0. 0. 0. 0. 0.]]

 [[0. 1. 1. 0. 0.]
  [0. 1. 1. 0. 0.]
  [0. 0. 0. 0. 0.]]

 [[0. 0. 1. 1. 0.]
  [0. 0. 1. 1. 0.]
  [0. 0. 0. 0. 0.]]

 [[0. 0. 0. 1. 1.]
  [0. 0. 0. 1. 1.]
  [0. 0. 0. 0. 0.]]

 [[0. 0. 0. 0. 0.]
  [1. 1. 0. 0. 0.]
  [1. 1. 0. 0. 0.]]

 [[0. 0. 0. 0. 0.]
  [0. 1. 1. 0. 0.]
  [0. 1. 1. 0. 0.]]

 [[0. 0. 0. 0. 0.]
  [0. 0. 1. 1. 0.]
  [0. 0. 1. 1. 0.]]

 [[0. 0. 0. 0. 0.]
  [0. 0. 0. 1. 1.]
  [0. 0. 0. 1. 1.]]]

I wish it can be faster so I create an array for index and trying to use np.vectorize. But as manual described, vectorize is not for performance. And my goal is running through an array with shape of (10^6, 15, 15) which end up with 10^6 iteration. I hope there are some cleaner solution can get rid of all the for-loop.

This is the first time I using stack overflow, any suggestion are appreciated.

Thank you.

Upvotes: 3

Views: 75

Answers (1)

B. M.
B. M.

Reputation: 18668

A efficient solution using numpy.lib.stride_tricks, which can "view" all the possibilities.

N=4 #tray size #(square)
P=3  # chunk size
R=N-P

from numpy.lib.stride_tricks import as_strided

tray = zeros((N,N),numpy.int32)
chunk = ones((P,P),numpy.int32)
tray[R:,R:] = chunk
tray = np.vstack((tray,tray))
view = as_strided(tray,shape=(R+1,R+1,N,N),strides=(4*N,4,4*N,4))
a_view = view.reshape(-1,N,N)
a_hard = a_view.copy()

Here is the result :

In [3]: a_view
Out[3]: 
array([[[0, 0, 0, 0],
        [0, 1, 1, 1],
        [0, 1, 1, 1],
        [0, 1, 1, 1]],

       [[0, 0, 0, 0],
        [1, 1, 1, 0],
        [1, 1, 1, 0],
        [1, 1, 1, 0]],

       [[0, 1, 1, 1],
        [0, 1, 1, 1],
        [0, 1, 1, 1],
        [0, 0, 0, 0]],

       [[1, 1, 1, 0],
        [1, 1, 1, 0],
        [1, 1, 1, 0],
        [0, 0, 0, 0]]])

a_view is just a view on possible positions of a chunk on the tray. It doesn't cost any computation, and it just uses twice the tray space. a_hard is a hard copy, necessary if you need to modify it.

Upvotes: 2

Related Questions