Получить словарь определенного вида

Рейтинг: 1Ответов: 1Опубликовано: 27.03.2023

Имеется датафрейм следующего вида:

q = [[nan, nan],
 [nan, nan],
 ['29', 'B'],
 ['29', 'D'],
 ['29', 'C'],
 [nan, nan],
 ['30', 'B'],
 ['30', 'D'],
 ['30', 'C'],
 [nan, nan],
 ['30', 'E'],
 ['30', 'G'],
 ['30', 'F'],
 [nan, nan],
 ['32', 'B'],
 ['32', 'D'],
 ['32', 'C'],
 [nan, nan],
 ['31', 'B'],
 ['31', 'D'],
 ['31', 'C']]

Названия колонок:

columns = ['source', 'own']

Требуется получить словарь следующего вида:

dict1 = {'29': ['B', 'D', 'C'],
 '30': ['B', 'D', 'C', '', 'E', 'G', 'F'],
 '32': ['B', 'D', 'C'],
 '31': ['B', 'D', 'C'],
 '33': ['B', 'D', 'C', '', 'E', 'G', 'F']}

На данный момент есть такое решение:

def get_correct_column_order(data_schema):
    correct_column_order = defaultdict(list)
    for i in range(data_schema.shape[0]):
        source = data_schema['source'].iloc[i]
        if isinstance(source, str):
            value = data_schema['own'].iloc[i]
            correct_column_order[source].append(value)
    return dict(correct_column_order)

df = pd.DataFrame(q, columns=columns)
get_correct_column_order(df)

и получаю словарь

dict2 = {'29': ['B', 'D', 'C'],
 '30': ['B', 'D', 'C', 'E', 'G', 'F'],
 '32': ['B', 'D', 'C'],
 '31': ['B', 'D', 'C'],
 '33': ['B', 'D', 'C', 'E', 'G', 'F']}
python словари dict dataframe
Источник: Stack Overflow на русском

Ответы

▲ 1Принят

Зачем так сложно, обычный groupby работает. Но чтобы пустая строка ещё была, придётся сначала пошаманить:

mask = df['source'].isnull() & (df['source'].ffill() == df['source'].bfill())
df.loc[mask, 'source'] = df['source'].ffill().loc[mask]
df.loc[mask, 'own'] = ''

df.groupby('source').agg(list)['own'].to_dict()

Вывод:

{'29': ['B', 'D', 'C'],
 '30': ['B', 'D', 'C', '', 'E', 'G', 'F'],
 '31': ['B', 'D', 'C'],
 '32': ['B', 'D', 'C']}

Если вы используете циклы в Pandas, значит вы что-то не то делаете, обычно всё можно решить "векторным" способом.