CODEWITHSUNDEEP
pythonnumpynumpy-slicing
vsocrates
vsocrates

Reputation: 184

Slice all except first element, unless single element

I'd like to slice a numpy array to get all but the first item, unless there's only one element, in which case, I only want to select that element (i.e. don't slice).

Is there a way to do this without using an if-statement?

x = np.array([1,2,3,4,5])
y = np.array([1])
print(x[1:]) # works
print(y[1 or None:]) # doesn't work

I tried the above, but it didn't work.

Upvotes: 2

Views: 1998

Answers (5)

Sangeerththan Balachandran
Sangeerththan Balachandran

Reputation: 593

Check the array size is greater than 1 if the case you can delete the first element from array and it will give new array without it.

print(np.delete(x, 0))

Now you can get a new array which will contain only remaining items other than first.

Upvotes: 0

joanis
joanis

Reputation: 12290

Use a ternary expression to make the logic clear but still be able to use it as a function argument or inside other expressions.

The ternary expression x if len(x) == 1 else x[1:] works and is very clear. And you can use it as a parameter in a function call, or in a larger expression.

E.g.:

>>> x = np.array([1,2,3,4,5])
>>> y = np.array([1])
>>> print(x if len(x) == 1 else x[1:])
[2 3 4 5]
>>> print(y if len(y) == 1 else y[1:])
[1]

Musings on other solutions

I'm not sure if you're looking for the most concise code possible to do this, or just for the ability to have the logic inside a single expression.

I don't recommend fancy slicing solutions with negative indexing, for the sake of legibility of your code. Think about future readers of your code, even yourself in a year or two.

Using this is a larger expression

In the comments, you mention you need a solution that can be incorporated into something like a comprehension. The ternary expression can be used as is within a comprehension. For example, this code works:

l = [np.array(range(i)) for i in range(5)]
l2 = [
    x if len(x) == 1 else x[1:]
    for x in l
]

I've added spacing to make the code easier to read, but it would also work as a one liner:

l2 = [x if len(x) == 1 else x[1:] for x in l]

EDIT note

Earlier, I thought you wanted the first element extracted from the list in the single-element case, i.e., x[0], but I believe you actually want that single-element list unsliced, i.e., x, so I've updated my answer accordingly.

Upvotes: 0

Mustafa Aydın
Mustafa Aydın

Reputation: 18315

A way to write that without a conditional is to use negative indexing with -len(arr) + 1:

>>> x = np.array([1,2,3,4,5])
>>> y = np.array([1])

>>> x[-len(x)+1:]
array([2, 3, 4, 5])

>>> y[-len(y)+1:]
array([1])

If the array has N elements where N > 1, slice becomes -N+1:. Since -N+1 < 0, it is effectively (N + (-N + 1)): === 1:, i.e, first one onwards.

Elsewhen N == 1, slice is 0:, i.e., take the first element onwards which is the only element.

Because of how slicing works, an empty array (i.e., N = 0 case) will result in an empty array too.

Upvotes: 2

MarkBeras
MarkBeras

Reputation: 624

Just move None out of the brackets

x = [1,2,3,4,5]
y = [1]
print(x[1:]) 
print(y[1:] or None)

Upvotes: 0

akuiper
akuiper

Reputation: 215137

You can just write an if / else:

x[1 if len(x) > 1 else 0:]
array([2, 3, 4, 5])

y[1 if len(y) > 1 else 0:]
array([1])

Or:

y[int(len(y) > 1):]
array([1])

x[int(len(x) > 1):]
array([2, 3, 4, 5])

Upvotes: 1

Related Questions