Index a NumPy array row-wise

Say I have a NumPy array:

>>> X = np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]])
>>> X
array([[ 1,  2,  3,  4],
       [ 5,  6,  7,  8],
       [ 9, 10, 11, 12]])

and an array of indexes that I want to select for each row:

>>> ixs = np.array([[1, 3], [0, 1], [1, 2]])
>>> ixs
array([[1, 3],
       [0, 1],
       [1, 2]])

How do I index the array X so that for every row in X I select the two indices specified in ixs?

So for this case, I want to select element 1 and 3 for the first row, element 0 and 1 for the second row, and so on. The output should be:

array([[2, 4],
       [5, 6],
       [10, 11]])

A slow solution would be something like this:

output = np.array([row[ix] for row, ix in zip(X, ixs)])

however this can get kinda slow for extremely long arrays. Is there a faster way to do this without a loop using NumPy?

EDIT: Some very approximate speed tests on a 2.5K * 1M array (10GB):

np.array([row[ix] for row, ix in zip(X, ixs)]) 0.16s

X[np.arange(len(ixs)), ixs.T].T 0.175s

X.take(idx+np.arange(0, X.shape[0]*X.shape[1], X.shape[1])[:,None]) 33s

np.fromiter((X[i, j] for i, row in enumerate(ixs) for j in row), dtype=X.dtype).reshape(ixs.shape) 2.4s

Solution:

You can use this:

X[np.arange(len(ixs)), ixs.T].T

Here is the reference for complex indexing.

Why was & used over && in java when comparing two bools?

I was looking through code in Guava https://github.com/google/guava and I see a lot of cool optimizations.

I was wondering if using & over && is an optimization and if it is, why is it? Could it be a style choice?

We are squaring an int b in the IntMath.checkedPow function. We want to check that b * b does not overflow:

checkNoOverflow(-FLOOR_SQRT_MAX_INT <= b & b <= FLOOR_SQRT_MAX_INT);
        b *= b;

In this example, why was & used over &&?

Edit: Matt is correct. I compiled this Java code in Java 8:

public static boolean and (boolean a, boolean b){
    return a && b;
}

public static boolean andBit (boolean a, boolean b){
    return a & b;
}

I looked at the Byte code using Intellij. I see that we are using branching in the “and” function because we have IFEQ and GOTO.

  // access flags 0x9
  public static and(ZZ)Z
   L0
    LINENUMBER 8 L0
    ILOAD 0
    IFEQ L1
    ILOAD 1
    IFEQ L1
    ICONST_1
    GOTO L2
   L1
   FRAME SAME
    ICONST_0
   L2
   FRAME SAME1 I
    IRETURN
   L3
    LOCALVARIABLE a Z L0 L3 0
    LOCALVARIABLE b Z L0 L3 1
    MAXSTACK = 1
    MAXLOCALS = 2

  // access flags 0x9
  public static andBit(ZZ)Z
   L0
    LINENUMBER 12 L0
    ILOAD 0
    ILOAD 1
    IAND
    IRETURN
   L1
    LOCALVARIABLE a Z L0 L1 0
    LOCALVARIABLE b Z L0 L1 1
    MAXSTACK = 2
    MAXLOCALS = 2
}

To answer the question, & is faster if the cost of evaluating an extra <= is faster than branching.

Erwin’s comment made me look closer at the actual while loop. b *=b is in a while loop which may repeat a lot. However, b can only be negative in the first loop because when we pass b *= b: b will be positive from then on.

Solution:

Since conditional branches are somewhat expensive, but comparing the values of integers is very fast, the fastest way to implement that expression in machine instructions would evaluate both sub-expressions and would not include a conditional branch.

The way it is written in your code snippet does reflect the best way to implement it. We usually expect the compiler to figure stuff like this out by itself, but maybe they tested and found that it did not…

Or maybe the developer just wrote it the faster way because the alternative would be to write it in a way that might be slower — and there’s certainly no good reason to do that!