CODEWITHSUNDEEP
pythonnumpy
Corei13
Corei13

Reputation: 261

Numpy: Efficient Way To Extract Subarray

I have a large two dimentional numpy array A of dimension around 250000 x 30 and two one dimensional numpy arrays x and y. I want to extract the sub-array of A with rows in x and columns in y. Which method is more efficient?

  1. A[x[:,np.newaxis], y]
  2. A[np.ix_(x,y)]

Upvotes: 0

Views: 711

Answers (1)

DrV
DrV

Reputation: 23550

Benchmark it!

import numpy as np

# some data
A = np.random.random((250000, 30))

# some random indices
x = np.random.randint(0, 250000, 150000)
y = np.random.randint(0, 30, 10)

def method1(A, x, y):
    return A[x[:, np.newaxis], y]

def method2(A, x, y):
    return A[np.ix_(x,y)]

def method3(A, x, y):
    return A[x][:,y]

def method4(A, x, y):
    return A[:,y][x]

These three methods give the following benchmarks:

method1: 87.7 ms
method2: 89.2 ms
method3: 115 ms
method4: 141 ms

So, the answer is that there is not real difference between the two methods in the question.

Upvotes: 2

Related Questions