labeling_1d.py

import numpy as np
import pandas as pd

def is_not_join(i):
    if i in ["A", "B", "C"]:
        ret = True
    else:
        ret = False
    return ret

s1 = pd.Series(["A", "B", "A", np.nan, "D", "E", np.nan, "A", "D", "E","A"]).values

lbl_num = 1
lbl = np.zeros_like(s1)
i = 0
first = True
while i < len(s1):
    if first:
        if pd.notnull(s1[i]):
            lbl_num = 1
            lbl[i] = lbl_num
        else:
            lbl_num = 0
        i += 1
        first = False
        continue
    if pd.isnull(s1[i-1]):
        if pd.isnull(s1[i]):
            pass
        else:
            lbl_num += 1
            lbl[i] = lbl_num
    else:
        if pd.isnull(s1[i]):
            pass
        else:
            if is_not_join(s1[i-1]):
                lbl_num += 1
                lbl[i] = lbl_num
            elif is_not_join(s1[i]):
                lbl_num += 1
                lbl[i] = lbl_num
            else:
                lbl[i] = lbl_num
    b = s1[i]
    i += 1

b = lbl