CODEWITHSUNDEEP
pythonnumpy
Vasile Alecsandri
Vasile Alecsandri

Reputation: 33

Numpy: How is this code different from each other?

Let r be an array where each element is an column index (less than N, size of r is M) and P be a MxN array.

The following two snippets behave differently. Why?

1.

P[:, r] += 1

2.

for i in range(len(r)):
    P[i, r[i]] += 1

Upvotes: 1

Views: 43

Answers (2)

Paul Panzer
Paul Panzer

Reputation: 53089

The first one selects an entire column for each element of r. The second just an element. You can directly compare the two cases like so:

>>> P = np.arange(12).reshape(4, 3)
>>> r = np.random.randint(0, 3, (4,))
>>> r
array([1, 1, 2, 0])
>>> 
>>> P[:, r]
array([[ 1,  1,  2,  0],
       [ 4,  4,  5,  3],
       [ 7,  7,  8,  6],
       [10, 10, 11,  9]])
>>> P[np.arange(4), r]
array([1, 4, 8, 9])

As you can see the second produces essentially the diagonal of the first.

You may profit from having a look at section "Combining advanced and basic indexing" in the numpy docs.

Upvotes: 3

Ryan Tam
Ryan Tam

Reputation: 855

P[:, r] in the first snippet selects the entire first axis (: is like 0:-1 here) and r-th second axis.

P[i, r[i]] in the for loop selects only the i-th first axis and r[i]-th second axis, which is just a single element.

Once that's clear, it's no surprising the two give different results.

Upvotes: 1

Related Questions