CODEWITHSUNDEEP
pythonnumpy
Shark Deng
Shark Deng

Reputation: 1068

slice an array without knowing its dimensions

I want to slice an array without knowing its dimensions. Indexes(start and end are given in list format. How to do this? Thank you.

Slice 1d array

import numpy as np
a = np.array([1, 2, 3, 4, 5, 6])
# idx: [2, 5)
print(a[2:5])
# [3, 4, 5]

Slice 2d array

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

start = [0, 1]
end = [1, 2]

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

Slice N-d array?

a = np.array([
    [[1, 0, 2],
    [2, 1, 0],
    [5, 6, 3]],
    
    [[2, 1, 3],
    [3, 2, 1],
    [1, 4, 6]]
])

start = [0, 0, 1]
end = [1, 1, 2]

result = a[start[0]:end[0], start[1]:end[1], start[2]:end[2]] # key point: how to adapt this?
print(result)

Upvotes: 0

Views: 303

Answers (1)

Kapocsi
Kapocsi

Reputation: 1022

If you have a list of start and end indexes, you can construct a tuple of slice objects and do the following to slice an array without knowing its dimensions:

a[tuple(slice(*indexes) for indexes in zip(start, end))]

Or:

a[tuple(slice(st, en) for st, en in zip(start, end))]

Or:

from itertools import starmap

a[tuple(starmap(slice, zip(start, end)))]

In action:

a = np.array([
    [[1, 0, 2],
    [2, 1, 0],
    [5, 6, 3]],
    
    [[2, 1, 3],
    [3, 2, 1],
    [1, 4, 6]]
])

start = [0, 0, 1]
end = [1, 1, 2]

result1 = a[start[0]:end[0], start[1]:end[1], start[2]:end[2]] # key point: how to adapt this?

result2 = a[tuple(slice(*indexes) for indexes in zip(start, end))] 

assert np.all(result2 == result2)

Upvotes: 1

Related Questions