Python - slice array until certain condition is met

Question

I need to slice an array from a given index until a certain condition is met.

>>> a = numpy.zeros((10), dtype='|S1')
>>> a[2] = 'A'
>>> a[4] = 'X'
>>> a[8] = 'B'
>>> a
array(['', '', 'A', '', 'X', '', '', '', 'B', ''], dtype='|S1')

For instance, for the above array I want a subset from a given index until first non-zero values in both directions. For example, for index values 2, 4, 8 the results would be:

['', '', A, '']      # 2
['', X, '', '', '']  # 4
['', '', '', B, '']  # 8

Any suggestions on the simplest way to do this using the numpy API? Learning python and numpy, would appreciate any help. Thanks!

Answer 1

If you set up your problem like this:

import numpy
a = numpy.zeros((10), dtype=str)
a[2] = 'A'
a[4] = 'X'
a[8] = 'B'

You can easily get the indices of non-empty strings like so:

i = numpy.where(a!='')[0]  # array([2, 4, 8])

Alternatively, numpy.argwhere(..) also works well.

Then you can slice away using this array:

out2 = a[:i[1]]        # 2   ['' '' 'A' '']
out4 = a[i[0]+1:i[2]]  # 4   ['' 'X' '' '' '']

etc.

Answer 2

this is a work for masked arrays, numpy.ma has lots of functions for working with subsets.

a = np.zeros((10), dtype=str)
a[2] = 'A'
a[4] = 'X'
a[8] = 'B'

let's mask out not empty elements:

am=np.ma.masked_where(a!='', a)

np.ma.notmasked_contiguous goes through the array (very efficiently) and finds all the slices of contiguous elements where the array is not masked:

slices = np.ma.notmasked_contiguous(am)
[slice(0, 1, None), slice(3, 3, None), slice(5, 7, None), slice(9, 9, None)]

so, the array is continuously empty between element 5 and 7 for example. Now you only have to join the slices you are interested in, first you get the starting index of each slice:

slices_start = np.array([s.start for s in slices])

then you get the location of the index you are looking for:

slices_start.searchsorted(4) #4
Out: 2

So you want slice 1 and 2: a[slices[1].start:slices[2].stop+1] array(['', 'X', '', '', ''], dtype='|S1')

or let's try 8:

i = slices_start.searchsorted(8)
a[slices[i-1].start:slices[i].stop+1]
Out: array(['', '', '', 'B', ''], 
  dtype='|S1')

If should probably play a bit with this in ipython for understanding it better.

Answer 3

Note that this could be cleanly done in pure python using itertools and functools.

import functools, itertools
arr = ['', '', 'A', '', 'X', '', '', '', 'B', '']

f = functools.partial(itertools.takewhile, lambda x: not x)
def g(a, i):
    return itertools.chain(f(reversed(a[:i])), [a[i]], f(a[i+1:]))

We define f as the sub-iterator found by looking until the element evaluates as true, and g as the combination of applying this on the reversed area of the list before the index and the list after the index.

This returns generators which can be casted to lists that contain our results.

>>> list(g(arr, 2))
['', '', 'A', '']
>>> list(g(arr, 4))
['', 'X', '', '', '']
>>> list(g(arr, 8))
['', '', '', 'B', '']