首页 > 解决方案 > Numpy:根据所选索引将部分数组替换为另一个完整数组

问题描述

我有两个 np 数组 A 和 B 以及一个索引列表,其中:

A = np.arange(12).reshape(3,2,2)
>>> A
array([[[ 0,  1],
        [ 2,  3]],

       [[ 4,  5],
        [ 6,  7]],

       [[ 8,  9],
        [10, 11]]])

>>> B:
array([[[10, 10],
        [10, 10]],

       [[20, 20],
        [20, 20]]])

index_list = [0,2]

我想根据 index_list 给出的索引将数组 A 替换为整个数组 B,索引为 0 和 2(对应于 A 在轴 = 0 处的索引):

# desired output:
array([[[ 10,  10],
        [ 10,  10]],

       [[ 4,  5],
        [ 6,  7]],

       [[ 20,  20],
        [20, 20]]])

我一直在考虑以更有效的方式实现这个过程,但只提出了两个循环:

row_count = 0
for i,e in enumerate(A):
    for idx in index_list:
        if i == idx:
            A[i,:,:] = B[row_count,:,:]
            row_count += 1

这是计算昂贵的,因为我有两个巨大的 np 数组要处理......如果有人知道如何以矢量化方式或更有效的方式实现这个过程,将不胜感激。谢谢!

标签: pythonarraysnumpyindexingnumpy-slicing

解决方案


您是否尝试过像这样分配:

A[index_list[0]] = B[0]
A[index_list[1]] = B[1]

我想如果你有更多的索引index_list,并且要替换更多,那么你可以创建一个循环。


推荐阅读