I have a numpy array that holds board game states for all possible moves, and I want to summarize some of those moves. I'm struggling to vectorize that code and avoid a for loop when I choose which moves I want to summarize.
Here's a simplified example of what I'm trying to do. I create a 3x3x3x3 array that represents the spaces that could be attacked by a queen at each board square of a 3x3 chess board. In other words, the first two dimensions are the coordinates of a queen on the board, and the last two dimensions are boolean flags of whether the queen could attack that square.
Then I select some squares, and count up how many of those squares could attack each square of the board. That counting step is what I'm trying to do without a Python for loop.
Here's the example code:
import numpy as np
size = 3
patterns = np.zeros((size, size, size, size), dtype=bool)
for i in range(size):
for j in range(size):
patterns[i, j, i, :] = True
patterns[i, j, :, j] = True
for i2 in range(size):
shift = i2 - i
j2 = j + shift
if 0 <= j2 < size:
patterns[i, j, i2, j2] = True
j3 = j - shift
if 0 <= j3 < size:
patterns[i, j, i2, j3] = True
active_positions = np.array([[0, 1, 0],
[1, 0, 0],
[0, 0, 0]], dtype=bool)
# This is the part I want to vectorize:
counts = np.zeros((size, size))
for i in range(size):
for j in range(size):
if active_positions[i, j]:
counts += patterns[i, j]
print(counts)
Is there a way to do that counting without using a for loop?
Use active_positions to mask patterns and sum along the zeroth (first) axis (that will be an element-wise sum of each field, which in this case is two 3x3 arrays).
counts = patterns[active_positions].sum(0)
Consider the intermediate array patterns[active_positions]; this has shape (2,3,3). Summing along the zeroth axis adds the two 3x3 arrays in an element-wise manner to produce the final (3,3) array that you're looking for.