python - 如何加速python中的循环

Question

我想加快这段代码的速度

      max_x=array([max(x[(id==dummy)]) for dummy in ids])

x 和 id 是相同维度的 numpy 数组，ids 是维度较小的数组。使用向量运算的快速方法是什么？

Answer 1

进一步向量化并不容易(据我所知)，除非 id 具有某种结构。否则瓶颈可能经常做 id==dummy ，但我能想到的唯一解决方案是使用排序，并且由于 np.max() 仍然缺乏减少功能需要相当多的 python 代码(编辑:实际上有一个通过 np.fmax 可用的 reduce 函数)。对于 x 为 1000x1000 且 id/ids 为 0..100 的情况，这大约快 3 倍，但由于它相当复杂，只有在具有许多 id 的较大问题时才值得这样做:

def max_at_ids(x, id, ids):
    # create a 1D view of x and id:
    r_x = x.ravel()
    r_id = id.ravel()
    sorter = np.argsort(r_id)

    # create new sorted arrays:
    r_id = r_id[sorter]; r_x = r_x[sorter]

    # unfortunatly there is no reduce functionality for np.max...

    ids = np.unique(ids) # create a sorted, unique copy, just in case

    # w gives the places where the sorted arrays id changes:
    w = np.where(r_id[:-1] != r_id[1:])[0] + 1

我最初提供的解决方案是在切片上执行纯 python 循环，但下面是一个更短(更快)的版本:

    # The result array:
    max_x = np.empty(len(ids), dtype=r_x.dtype)
    start_idx = 0; end_idx = w[0]
    i_ids = 0
    i_w = 0

    while i_ids < len(ids) and i_w < len(w) + 1:
        if ids[i_ids] == r_id[start_idx]:
            max_x[i_ids] = r_x[start_idx:end_idx].max()
            i_ids += 1
            i_w += 1
        elif ids[i_ids] > r_id[start_idx]:
            i_w += 1
        else:
            i_ids += 1
            continue # skip updating start_idx/end_idx

        start_idx = end_idx
        # Set it to None for the last slice (might be faster to do differently)
        end_idx = w[i_w] if i_w < len(w) else None

    return ids, max_x

编辑:计算每个切片最大值的改进版本:

有一种方法可以通过使用 np.fmax.reduceat 来删除 python 循环，如果切片很小(实际上非​​常优雅)，它可能会比前一个好很多:

# just to 0 at the start of w
# (or calculate first slice by hand and use out=... keyword argument to avoid even
# this copy.
w = np.concatenate(([0], w))
max_x = np.fmin.reduceat(r_x, w)
return ids, max_x

现在可能有一些小的东西可以让它更快一点。如果 id/ids 有一些结构，应该可以简化代码，并且可以使用不同的方法来实现更大的加速。否则这段代码的加速应该很大，只要有很多(唯一的)id(并且 x/id 数组不是很小)。请注意，代码强制执行 np.unique(ids)，但这可能是一个很好的假设。