simulate
simulate

Reputation: 1263

Selecting components of vectors in array

I have a tensor with shape (7, 2, 3) I want to select one of the two row vectors from each of the 7 2x3 matrices, i.e.

[
 [[0, 0, 0],
  [1, 1, 1]],

 [[0, 0, 0],
  [1, 1, 1]],
 ...x7
]

to

a = [
 [0, 0, 0],
 [0, 0, 0],
 [0, 0, 0]
 ...x7
]
b = [
 [1, 1, 1],
 [1, 1, 1],
 [1, 1, 1]
 ...x7
]

each with shape (7, 3).

How can I do this without reshape? (I find reshape to be kind of confusing when some dimensions are the same).

I also know of

np.array(map(lambda item: item[0], x)))

but I would like a more concise way if there is one.

Upvotes: 2

Views: 465

Answers (1)

anurag
anurag

Reputation: 1932

just use looped indexing: data[:, i, :] where i loops from 0 through 1

import numpy as np

a = np.array([
 [[0, 0, 0],
  [1, 1, 1]],

 [[0, 0, 0],
  [1, 1, 1]]
])
print(a[:, 1, :])

will produce

[[1 1 1]
 [1 1 1]]

Upvotes: 2

Related Questions