python - numpy(Python)中按组排列的数组元素的乘积

Question

我正在尝试构建一个返回数组元素子集的乘积的函数。基本上，我想构建一个执行此操作的 prod_by_group 函数:

values = np.array([1, 2, 3, 4, 5, 6])
groups = np.array([1, 1, 1, 2, 3, 3])

Vprods = prod_by_group(values, groups)

生成的 Vprods 应该是:

Vprods
array([6, 4, 30])

这里有一个关于元素总和的很好的答案，我认为它应该类似于: https://stackoverflow.com/a/4387453/1085691

我尝试先获取log，然后是sum_by_group，然后是exp，但遇到了数值问题。

对于按组划分的元素的最小值和最大值，这里还有一些其他类似的答案: https://stackoverflow.com/a/8623168/1085691

编辑:感谢您的快速回答!我正在试用它们。我应该补充一点，我希望它尽可能快(这就是我试图以某种矢量化方式在 numpy 中获取它的原因，就像我给出的示例一样)。

编辑:我评估了目前给出的所有答案，最好的答案是由下面的@seberg 给出的。这是我最终使用的完整功能:

def prod_by_group(values, groups):
    order = np.argsort(groups)
    groups = groups[order]
    values = values[order]
    group_changes = np.concatenate(([0], np.where(groups[:-1] != groups[1:])[0] + 1))
    return np.multiply.reduceat(values, group_changes)

Answer 1

如果您的组已经排序(如果没有，您可以使用 np.argsort 进行排序)，您可以使用 reduceat 功能对 进行排序ufunc(如果它们未排序，您必须先对它们进行排序才能高效地进行排序):

# you could do the group_changes somewhat faster if you care a lot
group_changes = np.concatenate(([0], np.where(groups[:-1] != groups[1:])[0] + 1))
Vprods = np.multiply.reduceat(values, group_changes)

如果您的小组很少，或者 mgilson 回答。但是如果你有很多组，那么这会更有效率。由于您避免为每个组的原始数组中的每个元素使用 bool 索引。此外，您还可以避免使用 reduceat 在 python 循环中进行切片。

当然，pandas 可以方便地完成这些操作。

编辑:抱歉，里面有 prod。 ufunc 是 multiply。您可以将此方法用于任何二进制 ufunc。这意味着它基本上适用于所有可以在两个输入数组上以元素方式工作的 numpy 函数。 (即乘法通常将两个数组按元素相乘，加法将它们相加，最大值/最小值等)