Pandas/NumPy: concisely label first N values matching a mask - pandas

Question

I have a sorted Series like this:
[2, 4, 5, 6, 8, 9]
I want to produce another Series or ndarray of the same length, where the first two odd numbers and the first two even numbers are labeled sequentially:
[0, 1, 2, _, _, 3]
The _ values I don't really care about. They can be zero.
Now I do it this way:
src = pd.Series([2, 4, 5, 6, 8, 9])
odd = src % 2 != 0
where = np.hstack((np.where(odd)[0][:2], np.where(~odd)[0][:2]))
where.sort() # maintain ordering - thanks to #hpaulj
res = np.zeros(len(src), int)
res[where] = np.arange(len(where))
Can you do it more concisely? The input will never be empty, but there might be no odds or no evens (in which case the result could have length 1, 2, or 3 instead of 4).

Answer 1

Great Problem! I'm still exploring and learning.
I've basically stuck with what you've done so far with modest tweaks for efficiency. I'll update if I think of anything else cool.
conclusions
So far, I've thrashed around alot and haven't improved much.
my answer
fast
odd = src.values % 2
even = 1 - odd
res = ((odd.cumsum() * odd) < 3) * ((even.cumsum() * even) < 3)
(res.cumsum() - 1) * res
alternative 1
pretty fast
a = src.values
odd = (a % 2).astype(bool)
rng = np.arange(len(a))
# same reason these are 2, we have 4 below
where = np.append(rng[~odd][:2], rng[odd][:2])
res = np.zeros(len(a), int)
# nature of the problem necessitates that this is 4
res[where] = np.arange(4)
alternative 2
not as quick, but creative
a = src.values
odd = a % 2
res = np.zeros(len(src), int)
b = np.arange(2)
c = b[:, None] == odd
res[(c.cumsum(1) * c <= 2).all(0)] = np.arange(4)
alternative 3
still slow
odd = src.values % 2
a = (odd[:, None] == [0, 1])
b = ((a.cumsum(0) * a) <= 2).all(1)
(b.cumsum() - 1) * b
timing
code
def pir3(src):
odd = src.values % 2
a = (odd[:, None] == [0, 1])
b = ((a.cumsum(0) * a) <= 2).all(1)
return (b.cumsum() - 1) * b
def pir0(src):
odd = src.values % 2
even = 1 - odd
res = ((odd.cumsum() * odd) < 3) * ((even.cumsum() * even) < 3)
return (res.cumsum() - 1) * res
def pir2(src):
a = src.values
odd = a % 2
res = np.zeros(len(src), int)
c = b[:, None] == odd
res[(c.cumsum(1) * c <= 2).all(0)] = np.arange(4)
return res
def pir1(src):
a = src.values
odd = (a % 2).astype(bool)
rng = np.arange(len(a))
where = np.append(rng[~odd][:2], rng[odd][:2])
res = np.zeros(len(a), int)
res[where] = np.arange(4)
return res
def john0(src):
odd = src % 2 == 0
where = np.hstack((np.where(odd)[0][:2], np.where(~odd)[0][:2]))
res = np.zeros(len(src), int)
res[where] = np.arange(len(where))
return res
def john1(src):
odd = src.values % 2 == 0
where = np.hstack((np.where(odd)[0][:2], np.where(~odd)[0][:2]))
res = np.zeros(len(src), int)
res[where] = np.arange(len(where))
return res
src = pd.Series([2, 4, 5, 6, 8, 9])
src = pd.Series([2, 4, 5, 6, 8, 9] * 10000)