See if index is contained within a slice object-CodePudding

I have an enum like this:

class AgeCategory(Enum):
    Child = slice(None, 16)
    YoungAdult = slice(16, 25)
    ...
    Old = slice(85, None)

It basically provides a range of ages (in years) which go in each category.

I would like a function to check which age range a certain value corresponds to. This was my idea:

def get_category(age: int) -> AgeCategory:
    for age_range in AgeCategory:
         if age in age_range.value #check if it's in the slice for that enum
             return age_range
    else:
        assert False, "unreachable"

However assert 5 in slice(1,10) fails. OFC I could do something like:

s: slice = age_range.value
if s.start <= age < s.stop: #bounds check
    return age_range

But that ignores the step argument and feels like reinventing the wheel a bit.

What's a pythonic way to express these age ranges? They are used as slices like this:

ya_data = np.sum(some_data[AgeCategory.YoungAdult])

CodePudding user response：

In your case it is not so much important, but this should be faster than iterating when working with large slices:

a_slice = slice(4, 15, 3)
def is_in_slice(a_slice, idx):
    if idx < a_slice.start or idx >= a_slice.stop:
        return False
    step = a_slice.step if a_slice.step else 1
    if (idx - a_slice.start) % step == 0:
        return True
    else:
        return False
    
test_array = np.where([is_in_slice(a_slice, idx) for idx in range(20)])[0]
print(test_array)

[ 4  7 10 13]

And test for very big slice:

a_big_slice = slice(0, 1_000_000_000_000, 5)
print(is_in_slice(a_big_slice, 999_000_000_005))
print(is_in_slice(a_big_slice, 999_000_000_004))

True
False

CodePudding user response：

With a sample slice:

In [321]: child=slice(None,16,4)

I was thinking of expanding it to a list or array. But arange can't handle the None:

In [323]: np.arange(child.start,child.stop,child.step)
Traceback (most recent call last):
  File "<ipython-input-323-b2d245f287ff>", line 1, in <module>
    np.arange(child.start,child.stop,child.step)
TypeError: unsupported operand type(s) for -: 'int' and 'NoneType'

np.r_ can. Here the numpy developers have gone to all the work of translating all the slice options:

In [324]: np.r_[child]
Out[324]: array([ 0,  4,  8, 12])
In [325]: 3 in _
Out[325]: False
In [327]: 4 in __
Out[327]: True

It may not be fastest, but appears to be most general purpose approach - with out a lot work on your part.