如果满足条件，如何交换 Nx2 numpy 数组行的元素？

Question

如果条件 swap[i] 为 True，我想将第 i 行的元素交换为 Nx2 numpy 数组 my_array

我的尝试：

def swap_positions_conditionally(my_array, swap):
    for i in range(np.shape(my_array)[0]):
        if swap[i]:
            my_array[i] = my_array[i][::-1]
    return my_array

工作正常，例如给出

my_array = array([[0, 1],
                  [2, 3],
                  [4, 5],
                  [6, 7],
                  [8, 9]])

swap = array([0, 0, 1, 1, 0])

产生预期结果

[[0 1]
 [2 3]
 [5 4]
 [7 6]
 [8 9]]

但是，可能有更惯用的表达方式来重写我的 swap_position_conditionally。
什么是更好（也更有效）的编写方式？

Answer 1

这是一个使用 np.take_along_axis:

np.take_along_axis(my_array, np.c_[swap, 1-swap], axis=1)

array([[0, 1],
       [2, 3],
       [5, 4],
       [7, 6],
       [8, 9]])

或基于布尔索引的一个：

swap = swap.astype(bool)
my_array[swap,:] = my_array[swap,::-1]

Answer 2

这里是直接使用 Numpy 的布尔索引：

import numpy as np

my_array = np.asarray([[0, 1],
                  [2, 3],
                  [4, 5],
                  [6, 7],
                  [8, 9]])

swap = np.array([0, 0, 1, 1, 0], dtype=bool)

my_array[swap, :] = my_array[swap,:][:,(1,0)]

分解重点线：

my_array[swap, :] = 表示 "Assign to the rows where swap is true"
my_array[swap,:] 表示 "select the whole row where swap is true"
[:,(1,0)] 表示 "for each row of what's to the left, swap the columns 0 and 1"

关于问题的 "more efficient" 部分...

所有测试的通用设置（种子确保序列相同）：

import timeit
setup= '''
import numpy as np

np.random.seed(42)
my_array = np.random.random([10000,2])
swap = np.random.random([10000]) > 0.5
'''

所有测试运行 1000 次迭代

原码：5.621秒

timeit.timeit('swap_positions_conditionally(my_array, swap)', setup=setup, number=1000)

如问题所示，将 swap_positions_conditionally 的定义添加到 setup。

这个答案：0.2657秒

timeit.timeit('my_array[swap, :] = my_array[swap,:][:,(1,0)]', setup=setup, number=1000)

Divakar 的回答：0.176 秒

timeit.timeit('np.where(swap[:,None]!=1,my_array,my_array[:,::-1])', setup=setup, number=1000)

牙兔第一个答案：0.214秒

timeit.timeit('np.take_along_axis(my_array, np.c_[swap, 1-swap], axis=1)', setup=setup, number=1000)

亚图的第二个答案：0.2547秒

timeit.timeit('my_array[swap,:] = my_array[swap,::-1]', setup=setup, number=1000)

结论

分析显示 Divakar 的版本是最快的。哪个更直观或可读性取决于口味，您可以选择您喜欢的那个（尽管我个人是索引符号可读性方面的粉丝......）

Answer 3

这是一种交换 Nx2 数组并在您尝试时使用负步长切片翻转更多列的方法 -

In [56]: np.where(swap[:,None]==1, my_array[:,::-1], my_array)
Out[56]: 
array([[0, 1],
       [2, 3],
       [5, 4],
       [7, 6],
       [8, 9]])

语法是：np.where(conditional_statement, choose_for_True, choose_for_False)。因此，在我们的例子中，当 swap 为 1 时，我们希望 flip/swap，否则不要。 [:,None] 部分需要在每一行中按元素执行此操作。如果swap已经是布尔数组，则跳过比较部分。

如果满足条件，如何交换 Nx2 numpy 数组行的元素？

How do I swap elements of rows of a Nx2 numpy array if a condition is satisfied?

python

swap

idioms

numpy