总结 numpy 中的更高维度

Question

我有一个 numpy 数组，用于保存所有可能的棋盘游戏状态，我想总结其中的一些棋步。当我选择要总结的动作时，我正在努力矢量化该代码并避免 for 循环。

这是我正在尝试做的事情的简化示例。我创建了一个 3x3x3x3 数组，表示 3x3 棋盘的每个棋盘方格上可能被皇后攻击的空间。换句话说，前两个维度是棋盘上皇后的坐标，最后两个维度是皇后是否可以攻击该方格的布尔标志。

然后我选择一些方格，并计算其中有多少个方格可以攻击棋盘上的每个方格。这个计数步骤就是我在没有 Python for 循环的情况下尝试做的。

这是示例代码：

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)

有没有一种方法可以在不使用 for 循环的情况下进行计数？

Answer 1

使用

active_positions

屏蔽

patterns

并沿第零（第一）轴求和。

counts = patterns[active_positions].sum(0)