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

Question

For all the folks who rock at vectorizing loops: I have two NumPy arrays of shape (N,) that contain indices to each other. Say we have a = np.asarray([0, 1, 2]) and b = np.array([1, 2, np.nan]) . The function should first look at a[0] to get 0 , then do b[0] to get 1 , then again a[1] to get 2 , and so on until we get np.nan . So the function is simply a[b[a[b[a[0]]]]] = np.nan . The output should contain two lists of values that were called for a and b respectively. Indices in b are always greater than in a , such that the process cannot get stuck.

I wrote a simple function that can do just this (wrapped with 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

The first list is lst[::2] :

[0, 2]

The second is lst[1::2] :

[2, 4]

Any way to vectorize this? Both arrays in inputs as well as both lists in output always have the same shape.

Answer 1

This is not a vectorized solution, but as a Numba solution it should be quite faster, and simpler. I changed the code slightly to use integers and -1 instead of np.nan , it is trivial to switch to this representation with something like b = np.where(np.isnan(b), -1, b) , and it makes the code more efficient. Instead of having a growing structure within the Numba function, I preallocate the output array in advance, so the loop can run much faster.

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]

Answer 2

Not vectorized, but here's recursive solution:

import numpy as np

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

def rec(i,b,a, a_out, b_out):
    if np.isnan(b[i]): return
    else:
        if not np.isnan(b[i]): a_out.append(i)
        rec(int(b[i]), a, b, b_out, a_out)
    return a_out, b_out

print(rec(0,b,a,[],[]))

Output

([0, 2], [2, 4])