使用 numpy 数组的切片索引

Question

Consider the following code:

import numpy as np
np.random.seed(0)

x = np.random.randint(0,9,(10,10,3))
x[np.array([0,1]):np.array([5,6])]

I get the following error:

TypeError: only integer scalar arrays can be converted to a scalar index

I guess I would need to unpack the arrays that index x but can't figure out a way to do it since * doesn't work.

The goal is to have x[0:5] and x[1:6]

EDIT :

Due to requests I'm posting here a solution for the problem I was looking for:

import numpy as np
np.random.seed(0)

x = np.random.randint(0,9,(10,10,3))
arr1,arr2 = np.array([0,1]),np.array([5,6])

slices = []
for i,j in zip(arr1,arr2):
    slices.append(x[i:j])

print(slices)

This way I can have both slices as requested ie x[0:5] and x[1:6] . If anyone has a method on how to approach this problem without a for loop it would be very helpful. Thanks.

Answer 1

here is the another one,

import numpy as np
np.random.seed(0)

x = np.random.randint(0,9,(10,10,3))
[x[i:j] for i,j in zip(np.array([0,1]),np.array([5,6]))]

Answer 2

if I have understood the question right, you should not use ":" in x[np.array([0,1]):np.array([0,1])]

use "," instead x[np.array([0,1]),np.array([0,1])]

Answer 3

You could generate an index array similar to this answer .

import numpy as np

x = np.random.randint(0,9,(10,10,3))
idcs = np.array([
    np.arange(0,5),
    np.arange(1,6),
])

print(x.shape)
# (10, 10, 3)
print(x[idcs].shape)
# (2, 5, 10, 3)