Search code examples
pythonnumpylist-comprehension

Value switcher for groups in Python/Numpy


I have a list:

groups = ['A', 'A', 'A', 'B', 'B', 'C', 'C', 'D']

I need to map each value to have an output like this, independently from the numbers of groups and elements inside:

[0,0,0,1,1,0,0,1]

The values in output should switch every time when the group is changing.


Solution

  • Using

    With a list comprehension and the walrus operator of python 3.8+:

    groups = ['A', 'A', 'A', 'B', 'B', 'C', 'C', 'D']
    
    flag = 0
    
    out = [flag if a==b else (flag:=1-flag) for a, b in zip(groups, groups[:1]+groups)]
    

    Or itertools:

    from itertools import groupby, chain
    
    out = list(chain.from_iterable([i%2]*len(list(g))
               for i, (_, g) in enumerate(groupby(groups))))
    

    Output:

    [0, 0, 0, 1, 1, 0, 0, 1]
    

    Using :

    import pandas as pd
    
    out = pd.factorize(groups)[0]%2
    

    Output:

    array([0, 0, 0, 1, 1, 0, 0, 1])
    

    Or:

    s = pd.Series(groups)
    out = (s.ne(s.shift(fill_value=s[0]))
           .cumsum().mod(2).tolist()
           )
    

    Output:

    [0, 0, 0, 1, 1, 0, 0, 1]
    

    Using :

    import numpy as np
    
    out = np.cumsum(groups != np.r_[groups[:1], groups[:-1]])%2
    

    Output:

    array([0, 0, 0, 1, 1, 0, 0, 1])