Question

我有一个名为test的非唯一原始数据。使用此输入，我想创建一个输出向量以及一组获得非零输出的rows和包含其输出的data。

import numpy as np

rows = np.array([3, 4])
test = np.array([1, 3, 3, 4, 5])
data = np.array([-1, 2])

我的预期输出是形状为test.shape的矢量。

output中的每个元素：

如果element位于rows且索引为i，output[i] = data[i]
否则，output[i] = 0

换句话说，以下内容会生成我的输出。

output = np.zeros(test.shape)
for i, val in enumerate(rows):
    output[test == val] = data[i]

有没有办法对此进行矢量化？

Answer 1

这是基于searchsorted -

的矢量化方法

# Get sorted index positions
idx = np.searchsorted(rows, test)

# Set out-of-bounds(invalid ones) to some dummy index, say 0
idx[idx==len(rows)] = 0

# Get invalid mask array found out by indexing data array
# with those indices and looking for matches
invalid_mask = rows[idx] != test

# Get data indexed array as output and set invalid places with 0s
out = data[idx]
out[invalid_mask] = 0

如果你挖一行，最后几行可以有两种选择 -

out = data[idx] * (rows[idx] == test) # skips using `invalid_mask`

out = np.where(invalid_mask, 0, data[idx])

Answer 2

这是一种仅在test和rows由不太大的整数组成时才有效的方法（非负数，但如果需要，可以放宽）。但那时速度很快：

>>> rows = np.array([3, 4])
>>> test = np.array([1, 3, 3, 4, 5])                                                                                        
>>> data = np.array([-1, 2])
>>> 
>>> limit = 1<<20
>>> assert all(a.dtype in map(np.dtype, np.sctypes['int']) for a in  (rows, test))
>>> assert np.all(rows>=0) and np.all(test>=0)
>>> mx = np.maximum(np.max(rows), np.max(test)) + 1
>>> assert mx <= limit
>>> lookup = np.empty((mx,), data.dtype)
>>> lookup[test] = 0
>>> lookup[rows] = data
>>> result = lookup[test]
>>> result
array([ 0, -1, -1,  2,  0])

矢量化这种非唯一键操作

2 个答案: