Панды группируют один и тот же столбец несколько раз, основываясь на разных значениях столбца.

Question

У меня есть пандас DataFrame, сгенерированный этим фрагментом:

elig = pd.DataFrame({'memberid': [1,1,1,1,1,1,2],
                     'monthid': [201711, 201712, 201801, 201805, 201806, 201807, 201810]})

, и я хотел бы выполнить операцию .groupby над memberid на основе непрерывных значений monthid, например,Мне бы хотелось, чтобы (очень) конечным результатом была таблица, похожая на эту:

memberid | start_month | end_month

    1    |    201711   |  201801
    1    |    201805   |  201807
    2    |    201810   |  201810

Мне было интересно, есть ли идиоматический способ Панд для этого.До сих пор я пробовал запутанный метод, определяющий new_elig = defaultdict(list), а затем внешнюю функцию:

def f(x):
    global new_elig
    new_elig[x.iloc[0]['memberid']].append(x.iloc[0]['monthid'])

и, наконец,

elig.groupby('memberid')[['memberid', 'monthid']].apply(f)

, который занимает около 5 минут для строк ~ 700k висходный DataFrame для создания new_elig, который затем я должен вручную проверять для каждого memberid, чтобы получить непрерывные диапазоны.

Есть ли лучший способ?Там должно быть одно: /

Answer 1

Вот один метод, который, я надеюсь, достаточно быстр для ваших нужд. это включает некоторую ручную арифметику по годам и месяцам. Это кажется грязным, но я думаю это быстрее, чем преобразование столбца monthid в серию Datetime с pd.to_datetime(elig['monthid'], format='%Y%m') и т. Д.

# Get the four-digit year with floor division

elig['year'] = elig['monthid']//100
elig['month'] = elig['monthid'] - elig['year']*100


# Boolean mask 1:
# If current row minus previous row is NOT 1 month, flag the row with True.
# Boolean mask 2:
# If months are contiguous (thus slipping past mask 1) 
# but memberid changes, flag the row with True.
# (This does not occur in your example data.)

mask1 = (elig['year']*12 + elig['month']).diff() != 1
mask2 = elig['memberid'] != elig['memberid'].shift()


# Convert the flag column to integer and take the cumulative sum.
# This converts the boolean flags into a column that assigns a 
# unique integer to each contiguous run of consecutive months belonging
# to the same memberid.

elig['run_id'] = (mask1 | mask2).astype(int).cumsum()

res = (
       elig.groupby('run_id')
           .agg({'memberid': 'first', 'monthid': ['first', 'last']})
           .reset_index(drop=True)
      )
res.columns = ['memberid', 'start_month', 'end_month']

res    
       memberid  start_month  end_month
    0         1       201711     201801
    1         1       201805     201807
    2         2       201810     201810