Expand integer ranges in pandas DataFrame column

Question

I have a dataframe that looks like:

d = {'value': ['a','b','c','d','e','f','g', 'h'],\
     'id'   : ['0101', '0208', '0103', '0405', '0105,0116,0117',
                '0108-0110', '0231, 0232, 0133-0150', '0155, 0152-0154, 0151']}
df = pd.DataFrame(d)
>>>
       value                     id
0      a                   0101
1      b                   0208
2      c                   0103
3      d                   0405
4      e                   0105
5      e                   0116
6      e                   0117
7      f                   0108
8      f                   0109
9      f                   0110
10     g  0231, 0232, 0133-0150
11     h  0155, 0152-0154, 0151

but I need to expand these IDs so that each row is a single number, so it looks more like:

   value    id
0      a  0101
1      b  0208
2      c  0103
3      d  0405
4      e  0105
5      e  0116
6      e  0117
7      f  0108
8      f  0109
9      f  0110
10     g   ...

where each row is duplicated where the IDs were grouped (with the ranges expanded, and leading zeros preserved for IDs less than 4 digits).

I've got as far as

df['id'].str.split(",")
df['id'].str.contains("-")

but I can't think of a good way to do this. Can anyone help?

Answer 1

You can write a little routine to flatten your ranges, and then repeat values from the original as necessary.

from itertools import chain

flattened = []
for x in df['id'].str.split(r',\s*'):
    flattened.append([])
    for y in x:
        if '-' in y:
            start, end = pd.to_numeric(y.split('-'))
            flattened[-1].extend(pd.RangeIndex(start, end+1))
        else: 
            flattened[-1].append(int(y))

repeats = [len(f) for f in flattened]

df_flat = pd.DataFrame({
        'value': df.value.repeat(repeats).values, 
        'id': list(chain.from_iterable(flattened))})
df_flat.tail(10)

   value   id
25     g  146
26     g  147
27     g  148
28     g  149
29     g  150
30     h  155
31     h  152
32     h  153
33     h  154
34     h  151

---

This turns out to be pretty performant, even for larger data.

df_ = df
df = pd.concat([df_] * 1000, ignore_index=True)

%timeit flatten(df)  # Function running code above.
244 ms ± 15.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Answer 2

Here's a way to do it:

s = (df['id'].str.split(r"[, ]|[-]")
             .apply(pd.Series)
             .stack()
             .reset_index(level=1, drop=True))

df.drop('id', axis =1).join(s.to_frame()).reset_index(drop=True)

 value     0
0      a  0101
1      b  0208
2      c  0103
3      d  0405
4      e  0105
5      e  0116
6      e  0117
7      f  0108
8      f  0109
9      f  0110
10     g  0231
11     g  0232
12     g  0133
13     g  0150
14     h  0155
15     h  0152
16     h  0154
17     h  0151