python - 有没有办法在 NumPy 中向量化两个指向彼此的数组？

Question

对于所有热衷于向量化循环的人:我有两个形状为 (N,) 的 NumPy 数组，它们包含彼此的索引。假设我们有 a = np.asarray([0, 1, 2]) 和 b = np.array([1, 2, np.nan])。该函数应首先查看 a[0] 以获取 0，然后执行 b[0] 以获取 1，然后再次 a[1] 得到 2，依此类推，直到我们得到 np.nan。所以这个函数就是a[b[a[b[a[0]]]]] = np.nan。输出应包含分别为 a 和 b 调用的两个值列表。 b 中的索引始终大于 a 中的索引，这样进程就不会卡住。

我写了一个简单的函数可以做到这一点(用 numba 封装 - 18.2 µs):

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

lst = []
while True:
    if len(lst) > 0:
        idx = lst[-1]
    else:
        idx = 0
    if len(lst) % 2 == 0:
        if idx < len(a) - 1:
            next_idx = a[idx]
            if np.isnan(next_idx):
                break
            lst.append(int(next_idx))
        else:
            break
    else:
        if idx < len(b) - 1:
            next_idx = b[idx]
            if np.isnan(next_idx):
                break
            lst.append(int(next_idx))
        else:
            break

第一个列表是lst[::2]:

[0, 2]

第二个是lst[1::2]:

[2, 4]

有什么方法可以对其进行矢量化？输入中的两个数组以及输出中的两个列表始终具有相同的形状。

Answer 1

这不是矢量化解决方案，但作为 Numba 解决方案，它应该更快、更简单。我稍微更改了代码以使用整数和 -1 而不是 np.nan，用 b = np.where 之类的东西切换到这种表示是微不足道的(np.isnan(b), -1, b)，它使代码更高效。我没有在 Numba 函数中使用不断增长的结构，而是提前预分配输出数组，因此循环可以运行得更快。

import numba as nb

def point_each_other(a, b):
    # Convert inputs to array if they are not already
    a = np.asarray(a)
    b = np.asarray(b)
    # Make output array in advance
    out = np.empty(len(a) + len(b), dtype=a.dtype)
    # Call Numba function
    n = point_each_other_nb(a, b, out)
    # Return relevan part of the output
    return out[:n]

@nb.njit
def point_each_other_nb(a, b, out):
    curr = 0
    i = 0
    while curr >= 0:
        # You can do bad input checking with the following
        # if i >= len(out):
        #     raise ValueError
        # Save current index
        out[i] = curr
        # Get the next index
        curr = a[curr]
        # Swap arrays
        a, b = b, a
        # Advance counter
        i += 1
    # Return number of stored indices
    return i - 1

# Test
a = np.array([0, 1, 2, 3, 4])
b = np.array([2, 3, 4, -1, -1])
out = point_each_other(a, b)
print(out[::2])
# [0 2 4]
print(out[1::2])
# [0 2]