NumPy 将函数应用于与另一个 numpy 数组对应的行组

Question

我有一个 NumPy 数组，每行代表一些 (x, y, z) 坐标，如下所示：

a = array([[0, 0, 1],
           [1, 1, 2],
           [4, 5, 1],
           [4, 5, 2]])

我还有另一个 NumPy 数组，其中包含该数组的 z 坐标的唯一值，如下所示：

b = array([1, 2])

如何将函数（我们称之为“f”）应用于 a 中与 b 中的值相对应的每个行组？例如，b 的第一个值是 1，因此我将获取 a 中 z 坐标为 1 的所有行。然后，我将一个函数应用于所有这些值。

最后，输出将是一个与 b 形状相同的数组。

我正在尝试对其进行矢量化以使其尽可能快。谢谢！

预期输出示例（假设 f 是 count()）：

c = array([2, 2])

因为数组 a 中有 2 行，数组 b 中的 z 值为 1，数组 a 中有 2 行，数组 b 中的 z 值为 2。

一个简单的解决方案是像这样迭代数组 b：

for val in b:
    apply function to a based on val
    append to an array c

我的尝试：

我尝试做这样的事情，但它只返回一个空数组。

func(a[a[:, 2]==b])

Answer 1

问题在于，具有相同 Z 的行组可以具有不同的大小，因此您无法将它们堆叠到一个 3D numpy 数组中，这将允许轻松地沿第三维应用函数。一种解决方案是使用 for 循环，另一种是使用np.split：

a = np.array([[0, 0, 1],
              [1, 1, 2],
              [4, 5, 1],
              [4, 5, 2],
              [4, 3, 1]])


a_sorted = a[a[:,2].argsort()]

inds = np.unique(a_sorted[:,2], return_index=True)[1]

a_split = np.split(a_sorted, inds)[1:]

# [array([[0, 0, 1],
#         [4, 5, 1],
#         [4, 3, 1]]),

#  array([[1, 1, 2],
#         [4, 5, 2]])]

f = np.sum  # example of a function

result = list(map(f, a_split))
# [19, 15]

但恕我直言，最好的解决方案是按照 FBruzzesi 的建议使用 pandas 和 groupby。然后您可以将结果转换为 numpy 数组。

编辑：为了完整起见，这里是其他两个解决方案

列表理解：

b = np.unique(a[:,2])
result = [f(a[a[:,2] == z]) for z in b]

熊猫：

df = pd.DataFrame(a, columns=list('XYZ'))
result = df.groupby(['Z']).apply(lambda x: f(x.values)).tolist()

这是我得到的性能图a = np.random.randint(0, 100, (n, 3))：

正如您所看到的，大约n = 10^5“分割解决方案”是最快的，但在那之后 pandas 解决方案的性能更好。