Home > Back-end >  Pandas: start a new group on every non-NA value
Pandas: start a new group on every non-NA value

Time:11-08

I am looking for a method to create an array of numbers to label groups, based on the value of the 'number' column. If it's possible?

With this abbreviated example DF:

number = [nan,nan,1,nan,nan,nan,2,nan,nan,3,nan,nan,nan,nan,nan,4,nan,nan]
df = pd.DataFrame(columns=['number'])
df = pd.DataFrame.assign(df, number=number)

Ideally I would like to make a new column, 'group', based on the int in column 'number' - so there would be effectively be array's of 1, ,2, 3, etc. FWIW, the DF is 1000's lines long, with sporadically placed int's.

The result would be a new column, something like this:

    number  group
0      NaN      0
1      NaN      0
2      1.0      1
3      NaN      1
4      NaN      1
5      NaN      1
6      2.0      2
7      NaN      2
8      NaN      2
9      3.0      3
10     NaN      3
11     NaN      3
12     NaN      3
13     NaN      3
14     NaN      3
15     4.0      4
16     NaN      4
17     NaN      4

All advice much appreciated!

CodePudding user response:

You can use notna combined with cumsum:

df['group'] = df['number'].notna().cumsum()

NB. if you had zeros: df['group'] = df['number'].ne(0).cumsum().

output:

    number  group
0      NaN      0
1      NaN      0
2      1.0      1
3      NaN      1
4      NaN      1
5      NaN      1
6      2.0      2
7      NaN      2
8      NaN      2
9      3.0      3
10     NaN      3
11     NaN      3
12     NaN      3
13     NaN      3
14     NaN      3
15     4.0      4
16     NaN      4
17     NaN      4

CodePudding user response:

You can use forward fill:

df['number'].ffill().fillna(0)

Output:

0     0.0
1     0.0
2     1.0
3     1.0
4     1.0
5     1.0
6     2.0
7     2.0
8     2.0
9     3.0
10    3.0
11    3.0
12    3.0
13    3.0
14    3.0
15    4.0
16    4.0
17    4.0
Name: number, dtype: float64
  • Related