Добавить 2-й массив в 3-й массив с постоянно меняющимся индексом быстро

Question

Я пытаюсь добавить 2-й массив в 3-й массив с постоянно меняющимся индексом, я придумываю следующий код:

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)

, который даст именно то, что я хочу:

[[[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.]]]

Хотелось бы, чтобы это было быстрее, поэтому я создаю массив для индекса и пытаюсь использовать np.vectorize. Но, как описано в руководстве, векторизация не для производительности. И моя цель состоит в том, чтобы пройти через массив с формой (10 ^ 6, 15, 15), который заканчивается 10 ^ 6 итераций. Я надеюсь, что есть какое-то более чистое решение, которое поможет избавиться от всей петли for.

Это первый раз, когда я использую переполнение стека, любые предложения приветствуются.

Спасибо.

Answer 1

Эффективное решение с использованием numpy.lib.stride_tricks , которое может «просматривать» все возможности.

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

Вот результат:

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 - это просто представление о возможных положениях блока в лотке. Это не требует каких-либо вычислений, а просто занимает вдвое больше места в трее. a_hard - это печатная копия, необходимая, если вам нужно ее изменить.