Question

我正在尝试优化以下代码，可能是通过在Cython中重写它：它只需要一个低维但相对较长的numpy数组，查看其列的0值，并在数组中将它们标记为-1 。代码是：

import numpy as np

def get_data():
    data = np.array([[1,5,1]] * 5000 + [[1,0,5]] * 5000 + [[0,0,0]] * 5000)
    return data

def get_cols(K):
    cols = np.array([2] * K)
    return cols

def test_nonzero(data):
    K = len(data)
    result = np.array([1] * K)
    # Index into columns of data
    cols = get_cols(K)
    # Mark zero points with -1
    idx = np.nonzero(data[np.arange(K), cols] == 0)[0]
    result[idx] = -1

import time
t_start = time.time()
data = get_data()
for n in range(5000):
    test_nonzero(data)
t_end = time.time()
print (t_end - t_start)

data是数据。 cols是要查找非零值的数据列数组（为简单起见，我将它设置为相同的列）。目标是计算一个numpy数组result，其中感兴趣的列非零的每一行都有1个值，而感兴趣的相应列的行为-1则为-1。 / p>

在15,000行3列的不太大的数组上运行此函数需要大约20秒。有没有办法可以加速？似乎大多数工作都是寻找非零元素并使用索引检索它们（调用nonzero并随后使用其索引。）这可以优化还是这是最好的？ Cython实现如何才能加快速度？

Answer 1

cols = np.array([2] * K)

这将非常缓慢。这会创建一个非常大的python列表，然后将其转换为numpy数组。相反，做一些像：

cols = np.ones(K, int)*2

那会更快

result = np.array([1] * K)

你应该这样做：

result = np.ones(K, int)

这会直接产生numpy数组。

idx = np.nonzero(data[np.arange(K), cols] == 0)[0]
result[idx] = -1

cols是一个数组，但是你可以传递一个2.此外，使用非零值会增加一个额外的步骤。

idx = data[np.arange(K), 2] == 0
result[idx] = -1

应该有同样的效果。

在Python中优化索引和检索numpy数组中的元素？

1 个答案: