Помогите оптимизировать код

Есть датафрейм вида введите сюда описание изображения

Необходимо получить два списка:

['Растения', 'Животные', 'Грибы']

и

[[{'Название': 'Роза', 'Количество': 15}, 
  {'Название': 'Ель', 'Количество': 2}], 
 [{'Название': 'Слон', 'Количество': 1}, 
  {'Название': 'Тигр', 'Количество': 3}],
 [{'Название': 'Подосиновик', 'Количество': 19}, 
  {'Название': 'Мухомор', 'Количество': 2}]]

Сейчас получаю эти данные с помощью следующего кода

def foo(df):
    tsarstva = []
    all = []
    one = []
    for i in range(len(df)):
        k = df.iloc[i]
        if k[0] not in tsarstva:
            tsarstva.append(k[0])
            if one:
                one_df = pd.DataFrame(one)
                one_df = one_df[
                    ['Название',
                     'Количество']
                ].copy()
                res = one_df.to_dict('records')
                all.append(res)
            one.clear()
        else:
            one.append(k)
    one_df = pd.DataFrame(one)
    one_df = one_df[
                    ['Название',
                     'Количество']
                ].copy()
    res = one_df.to_dict('records')
    all.append(res)
    return tsarstva, all

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


Ответы (1 шт):

Автор решения: strawdog

очень много лишнего, вы не используете функционал pandas

при исходном df:

     Царство     Название  Количество
0   Растения          NaN         NaN
1        NaN         Роза        15.0
2        NaN          Ель         2.0
3        NaN          NaN         NaN
4   Животные          NaN         NaN
5        NaN         Слон         1.0
6        NaN         Тигр         3.0
7        NaN          NaN         NaN
8      Грибы          NaN         NaN
9        NaN  Подосиновик        19.0
10       NaN      Мухомор         2.0
df["Царство"] = df["Царство"].fillna(method="ffill")
df = df.dropna()
names = df["Царство"].unique()
print(names)

['Растения' 'Животные' 'Грибы']

res = df.groupby("Царство")[["Название","Количество"]].apply(
    lambda x: x.to_dict("records")).values.tolist()

res:

[[{'Название': 'Подосиновик', 'Количество': 19.0},
  {'Название': 'Мухомор', 'Количество': 2.0}],
 [{'Название': 'Слон', 'Количество': 1.0},
  {'Название': 'Тигр', 'Количество': 3.0}],
 [{'Название': 'Роза', 'Количество': 15.0},
  {'Название': 'Ель', 'Количество': 2.0}]]
→ Ссылка