Оптимизация кода (цикл, условия)

Question

У меня есть google таблица, которая, к примеру, выглядит следующим образом:

далее через DataFrame библиотеки Pandas я получаю данные из этой таблицы для дальнейшего анализа. Есть некий Архив и есть связка Номер - Цвет. Задача - подсчитать количество каждого цвета в столбце D, но с условием, что номера из столбца C могут дублироваться столбце A, в таком случае количество Красного должно просуммироваться с Разноцветным. Дублироваться могут только номера, которым сопоставлен Красный цвет. (если это значение дублируется, то считать этот цвет не Красным, а Разноцветным и суммировать соответственно количество таких дублей не с Красными, а с Разноцветными).

Касаемо кода. Вот, что я смог придумать. Он вроде полностью отрабатывает задачу, но выглядит довольно грязно и слабочитаем. Можете подсказать как его причесать и возможно избавиться от многовложенности условий.

import pandas as pd
from sheets import *

# Столбец A (архив)
archive = pd.DataFrame(test_sheet.get_values('A2:A21'), columns=['Номер'])
# Столбцы C,D (номер, цвет)
column = pd.DataFrame(test_sheet.get_values('C2:D21'), columns=['Номер', 'Цвет'])

color = {
    'Красный': 'Красный|красный',
    'Белый': 'Белый|белый',
    'Синий': 'Синий|синий',
    'Зеленый': 'Зеленый|зеленый',
    'Разноцветный': 'Разноцветный|разноцветный',
    'Оранжевый': 'Оранжевый|оранжевый'
}

to_archive = 0  # Перенос в архив через данную переменную

for key in color.keys():
    col_ = column[column['Цвет'].str.contains(color[key])]
    if key == 'Красный' and len(col_) > 0:
        double = pd.merge(archive, col_, how='inner')
        if len(double) > 0:
            to_archive = len(double)
            print(f'{key}: {len(col_) - to_archive}')
        else:
            print(f'{key}: {len(col_) - len(double)}')
    elif key == 'Разноцветный':
        print(f'{key}: {len(col_) + to_archive}')
    else:
        print(f'{key}: {len(col_)}')
print('******************')

Вывод:

Красный: 3
Белый: 2
Синий: 3
Зеленый: 3
Разноцветный: 5
Оранжевый: 4
******************

Добавляю ссылку на файл с исходными данными

Answer 1

Вы совсем не используете инструментарий pandas. Ваша задача решается в полторы строки:

import pandas as pd

df = pd.read_csv("test.csv")

df.loc[df["Номер"].isin(df["Архив"]), "Цвет"] = "Разноцветный"
res = df.value_counts("Цвет")

получаете res:

Цвет
Разноцветный    5
Оранжевый       4
Зеленый         3
Красный         3
Синий           3
Белый           2
dtype: int64

БЛОГ НА HUSL

Оптимизация кода (цикл, условия)

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