Numpy - Use values as index of another arrays

user_5077986 0
Viewed 268

I face a problem with Numpy, I try to use the values of each row (of B) as the indexes of another multidimensional array (A):

>>> A
array([[0, 1, 2],
       [3, 4, 5],
       [6, 7, 8]])
>>> B
    array([[ 0,  1,  2,  3],
           [ 1,  5,  6,  7],
           [ 8,  9, 10, 11]])
>>> np.clip(B, 0, 2)
    array([[0, 1, 2, 2],
           [1, 2, 2, 2],
           [2, 2, 2, 2]])

Here is the expected result:

array([[0, 1, 2, 2],
       [4, 5, 5, 5],
       [8, 8, 8, 8]])

Any idea ? Thanks a lot for your help.

multidimensional-array indexing python-3.x numpy python
2 Answers

It seems that you are looking for np.take_along_axis:

A = np.array([
       [0, 1, 2],
       [3, 4, 5],
       [6, 7, 8]]
)
B = np.array([
           [ 0,  1,  2,  3],
           [ 1,  5,  6,  7],
           [ 8,  9, 10, 11]]
)
C = np.clip(B, 0, 2)

res = np.take_along_axis(A, C, 1)
print(res)
user_45859630

Try this:

result = np.clip(B, 0, 2)
for i, row in enumerate(result):
    for j, elem in enumerate(row):
        result[i][j] = A[i][elem]

After @hilberts_drinking_problem's answer, I have done some benchmarks using timeit. The answers, when tested with 10,000,000 loops took 103s and 82s, the latter being mine. Therefore, if your data is small but if you are going to be looping over it quite a lot, I would recommend you to use my answer. If such is not the case, I believe the convenience of hilberts_drinking_problem's answer makes it superior. For more information, see the comments to this answer.

Note that in both cases, I tested just the parts after clipping, meaning the preceding preparation part does not interfere with the benchmark.

user_84030700
Related