使用'None'进行成对运算的numpy索引

Question

If I had two numpy arrays that looked like this

a = np.array([1, 2])
b = np.array([3, 4])

and I wanted to add all pairwise combinations, I could easily do

c = a + b[:, None]
c
array([[4, 5],
       [5, 6]])

to get the result of 1+3 , 2+3 and 1+4 , 2+4 .

Why does this work? What is 'None' doing? I can print out

b[:, None]
[[3]
 [4]]

But I'm not sure why that tells numpy to do pairwise combos. I'm also curious about if it's efficiently implemented under the hood compared to, say, itertools.combinations.

Answer 1

To answer the first part of your question, b[:, None] is a special type of slicing that has identical behavior to b[:, np.newaxis] , in that it adds an axis of length 1 to your array.

>>> b.shape
(2,)
>>> b[:, None].shape
(2, 1)

This behavior is documented in the numpy docs ^[1] , emphasis mine:

The newaxis object can be used in all slicing operations to create an axis of length one. newaxis is an alias for None , and None can be used in place of this with the same result.

So now we have two arrays:

array([1, 2]) + array([[3],
                       [4]])

Summing these two arrays results in:

array([[4, 5],
       [5, 6]])

The "magic" behind this is numpy broadcasting ^[2] . This article ^[3] is an excellent resource for beginning to understand the topic.

---

The main takeaways from the article are as follows:

numpy operations are usually done element-by-element which requires two arrays to have exactly the same shape. However, this constraint is relaxed if both arrays have the same trailing axis, or if one of the trailing axes is equal to one (which is the behavior being exhibited in your case).

In your case, broadcasting occurs, so the operation is equivalent to summing the following 2x2 arrays:

array([[1, 2],    +  array([[3, 3],
       [1, 2]])            [4, 4]])

Which, since numpy operations are done element-by-element, will produce the desired output of:

array([[4, 5],
       [5, 6]])

---

_{^[1] https://docs.scipy.org/doc/numpy/reference/arrays.indexing.html#numpy.newaxis}

_{^[2] https://docs.scipy.org/doc/numpy-1.13.0/user/basics.broadcasting.html}

_{^[3] http://scipy.github.io/old-wiki/pages/EricsBroadcastingDoc}

Answer 2

What you're doing is adding a new axis: the "column"-axis. b didn't have that before, so now it's a column vector and would be added column wise; it will essentially act as if it was repeated in the columns, and a repeated in the rows:

a+b[:, None] = [1,2] + [[3], = [[1,2], + [[3],[3],
                        [4]]    [1,2]]    [4],[4]]

And here is how/why:

First things first: numpy does elementwise addition and multiplication by default. That means if a=np.array([1,2]) then a+2=np.array([1+2,2+2])=np.array([3,5]) .

import numpy as np
A = np.array([[1, 2],
              [3, 4]])
B = np.array([[1,1],
              [0,0]])

Here A+B would be element wise and we would get

A+B = [[1+1,2+1], = [[2,3],
       [3+0,4+0]] =  [3,4]]

If we transpose the matrix B (using BT )

now BT would have the value

B.T= np.array([[1,0],
               [1,0]])

And if we do it elementwise this time we'd get:

A+B.T=[[1, 2]  + [[1,0]  =  [[2,2],
       [3, 4]]    [1,0]]     [4,4]]

Another thing to notice is that ( B not transposed)

a = np.array([1,2])
B+a = [[1,2], + [[2, 3],
       [1,2]]    [1, 2]]

And it's also elementwise, but on both rows! That is a was practically "repeated" to have two rows, and added elementwise to B .

Further, Numpy docs says that None , in slicing, is a different way of writing np.newaxis . In your case, what the None option in slicing does is basically transposing the b -vector before addition!

The exact same result could have been obtained by

import numpy as np
a = np.array([1, 2])
b = np.array([3, 4])

c=a+b.reshape(2,1)
d=a+b.reshape(-1,1)
e=a+b[:, np.newaxis]

Here c , d and e have the same values!