CODEWITHSUNDEEP
pythonnumpyswapidioms
Davide Fiocco
Davide Fiocco

Reputation: 5914

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

I'd like to swap elements of the i-th row a Nx2 numpy array my_array if a condition swap[i] is True.

My attempt:

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

works fine, e.g. given

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

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

yields the expected result

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

However, there's likely a more idiomatic expression to rewrite my swap_position_conditionally.
What would be a better (and more efficient) way to write it?

Upvotes: 0

Views: 669

Answers (3)

GPhilo
GPhilo

Reputation: 19123

And here's one using Numpy's boolean indexing directly:

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)]

Breaking down the key line:

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

About the "more efficient" part of the question...

Common setup for all tests (seed ensures sequences are identical):

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
'''

All tests run for 1000 iterations

Original code: 5.621 seconds

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

Added the definition of swap_positions_conditionally to setup as shown in the question.

This answer: 0.2657 seconds

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

Divakar's answer: 0.176 seconds

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

Yatu's first answer: 0.214 seconds

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

Yatu's second answer: 0.2547 seconds

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

Conclusions

Profiling shows Divakar's version is the fastest. Whichever is more intuitive or readable is a matter of flavors, you can pick the one you prefer (I personally am a fan of the indexing notation readability-wise though...)

Upvotes: 1

yatu
yatu

Reputation: 88236

Here's one using 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]])

Or a boolean indexing based one:

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

Upvotes: 1

Divakar
Divakar

Reputation: 221574

Here's one way that swaps for Nx2 array and flips for more number of columns with the negative stepsize slicing as you were attempting -

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]])

Syntax is : np.where(conditional_statement, choose_for_True, choose_for_False). So, in our case, we want to flip/swap when swap is 1, else don't. That [:,None] part is needed to do this elementwise across each row. If swap is already a boolean array, skip the comparison part.

Upvotes: 1

Related Questions