Categorização baseada em prioridade usando pandas / python

Question

Tenho fatura e dados de código nos Dataframes abaixo

Faturas

df = pd.DataFrame({
    'invoice':[1,1,2,2,2,3,3,3,4,4,4,5,5,6,6,6,7],
    'code':[101,104,105,101,106,106,104,101,104,105,111,109,111,110,101,114,112],
    'qty':[2,1,1,3,2,4,7,1,1,1,1,4,2,1,2,2,1]
})

+---------+------+-----+
| invoice | code | qty |
+---------+------+-----+
|    1    |  101 |  2  |
+---------+------+-----+
|    1    |  104 |  1  |
+---------+------+-----+
|    2    |  105 |  1  |
+---------+------+-----+
|    2    |  101 |  3  |
+---------+------+-----+
|    2    |  106 |  2  |
+---------+------+-----+
|    3    |  106 |  4  |
+---------+------+-----+
|    3    |  104 |  7  |
+---------+------+-----+
|    3    |  101 |  1  |
+---------+------+-----+
|    4    |  104 |  1  |
+---------+------+-----+
|    4    |  105 |  1  |
+---------+------+-----+
|    4    |  111 |  1  |
+---------+------+-----+
|    5    |  109 |  4  |
+---------+------+-----+
|    5    |  111 |  2  |
+---------+------+-----+
|    6    |  110 |  1  |
+---------+------+-----+
|    6    |  101 |  2  |
+---------+------+-----+
|    6    |  114 |  2  |
+---------+------+-----+
|    7    |  112 |  1  |
+---------+------+-----+

Códigos

Hot =  [103,109]
Juice =  [104,105]
Milk =  [106,107,108]
Dessert =  [110,111]

Minha tarefa é adicionar uma coluna agora, com categorybase nas seguintes prioridades:

1. Se alguma fatura tiver mais de \$10\$ qtydeve ser categorizado como "Mega".
   Por exemplo, o total qtyda fatura 3 é \$12\$- \$4 + 7 + 1\$.

2. Se qualquer um dos invoice's codes estão na lista de leite ; a categoria deve ser "Saudável".
   Por exemplo, a fatura 2 contém o código 106 que está na lista de leite. Portanto, toda a fatura é categorizada Healthyindependentemente de outros itens.

3. Se algum dos invoices' codes estiver na lista de sucos ;

   1. Se o total qtyde sucos for igual a 1; a categoria deve ser "OneJuice".
      Por exemplo, a fatura 1 tem code104 e qty1.

   2. De outra forma; a categoria deve ser "ManyJuice".
      Por exemplo, a fatura 4 tem code104 e 105 com um total qtyde 2 - \$1 + 1\$.

4. Se qualquer um dos invoices's codes estão na lista quente ; a categoria deve ser "HotLovers".

5. Se qualquer um dos invoices's codes estão na lista de sobremesa ; a categoria deve ser "DessertLovers".

6. Todas as outras faturas devem ser categorizadas como "Outros".

Minha saída desejada é a seguinte.

+---------+------+-----+---------------+
| invoice | code | qty |    category   |
+---------+------+-----+---------------+
|    1    |  101 |  2  |    OneJuice   |
+---------+------+-----+---------------+
|    1    |  104 |  1  |    OneJuice   |
+---------+------+-----+---------------+
|    2    |  105 |  1  |    Healthy    |
+---------+------+-----+---------------+
|    2    |  101 |  3  |    Healthy    |
+---------+------+-----+---------------+
|    2    |  106 |  2  |    Healthy    |
+---------+------+-----+---------------+
|    3    |  106 |  4  |      Mega     |
+---------+------+-----+---------------+
|    3    |  104 |  7  |      Mega     |
+---------+------+-----+---------------+
|    3    |  101 |  1  |      Mega     |
+---------+------+-----+---------------+
|    4    |  104 |  1  |   ManyJuice   |
+---------+------+-----+---------------+
|    4    |  105 |  1  |   ManyJuice   |
+---------+------+-----+---------------+
|    4    |  111 |  1  |   ManyJuice   |
+---------+------+-----+---------------+
|    5    |  109 |  4  |   HotLovers   |
+---------+------+-----+---------------+
|    5    |  111 |  2  |   HotLovers   |
+---------+------+-----+---------------+
|    6    |  110 |  1  | DessertLovers |
+---------+------+-----+---------------+
|    6    |  101 |  2  | DessertLovers |
+---------+------+-----+---------------+
|    6    |  114 |  2  | DessertLovers |
+---------+------+-----+---------------+
|    7    |  112 |  1  |     Others    |
+---------+------+-----+---------------+

Eu tenho o seguinte. Funciona, mas parece muito ingênuo e nada Pythônico.
Quando eu o aplico ao conjunto de dados original, o código também fica muito lento.

# Calculating Priority No.1 
L = df.groupby(['invoice'])['qty'].transform('sum') >= 10
df_Large = df[L]['invoice'].to_frame()
df_Large['category'] = 'Mega'
df_Large.drop_duplicates(['invoice'], inplace=True)


# Calculating Priority No.2
df_1 = df[~L] # removing Priority No.1 calculated above
M = (df_1['code'].isin(Milk)
.groupby(df_1['invoice'])
.transform('any'))
df_Milk = df_1[M]['invoice'].to_frame()
df_Milk['category'] = 'Healthy'
df_Milk.drop_duplicates(['invoice'], inplace=True)

# Calculating Priority No.3

# 3.a Part -1

df_2 = df[~L & ~M]  # removing Priority No.1 & 2 calculated above
J_1 = (df_2['code'].isin(Juice)
.groupby(df_2['invoice'])
.transform('sum') == 1)
df_SM = df_2[J_1]['invoice'].to_frame()
df_SM['category'] = 'OneJuice'
df_SM.drop_duplicates(['invoice'], inplace=True)


# 3.b Part -2
J_2 = (df_2['code'].isin(Juice)
.groupby(df_2['invoice'])
.transform('sum') > 1)
df_MM = df_2[J_2]['invoice'].to_frame()
df_MM['category'] = 'ManyJuice'
df_MM.drop_duplicates(['invoice'], inplace=True)


# Calculating Priority No.4
df_3 = df[~L & ~M & ~J_1 & ~J_2]  # removing Priority No.1, 2 & 3 (a & b) calculated above
H = (df_3['code'].isin(Hot)
.groupby(df_3['invoice'])
.transform('any'))
df_Hot = df_3[H]['invoice'].to_frame()
df_Hot['category'] = 'HotLovers'
df_Hot.drop_duplicates(['invoice'], inplace=True)


# Calculating Priority No.5
df_4 = df[~L & ~M & ~J_1 & ~J_2 & ~H ] # removing Priority No.1, 2, 3 (a & b) and 4 calculated above
D = (df_4['code'].isin(Dessert)
.groupby(df_4['invoice'])
.transform('any'))
df_Dessert = df_4[D]['invoice'].to_frame()
df_Dessert['category'] = 'DessertLovers'
df_Dessert.drop_duplicates(['invoice'], inplace=True)

# merge all dfs
category = pd.concat([df_Large,df_Milk,df_SM,df_MM,df_Hot,df_Dessert], axis=0,sort=False, ignore_index=True)

# Final merge to the original dataset
df = df.merge(category,on='invoice', how='left').fillna(value='Others')

Answer 1

Seu código é bastante impressionante. Muitos programadores de python não sabem usar o pandas tão bem quanto você. Seu código pode não parecer muito "Pythônico", mas você fez um ótimo trabalho utilizando métodos vetorizados com indexação. Nesta resposta, incluo uma seção sobre convenções de código Python e uma segunda tentativa de otimizar seu código.

Convenções de código Python

Muitas empresas têm guias de estilo padronizados que tornam o código mais fácil de ler. Isso é inestimável quando muitas pessoas escrevem na mesma base de código. Sem consistência, o repo se degradaria em uma confusão de idiossincrasias.

Você deve considerar a adoção das seguintes convenções de código para tornar seu código mais fácil de ler:

1. Siga as convenções de nomenclatura de variável padrão: Guia de estilo do Google Python sobre nomenclatura
2. Inclua um espaço após as vírgulas: Guia de estilo do Google Python em espaços

# most python programmers use CaseLikeThis (pascal case) for class names
# constants are often written in CASE_LIKE_THIS (snake case)
SODA =  [101, 102]
HOT =  [103, 109]
JUICE =  [104, 105] # remember spaces after commas
MILK =  [106, 107, 108]
DESSERT =  [110, 111]

Tentar otimizar

Para otimizar seu código, você deve cronometrar quanto tempo leva cada etapa. Isso pode ser feito verificando o relógio antes e depois de um segmento de código.

import time

t0 = time.time() # check clock before (milliseconds elapsed since jan 1, 1970)
# segment you want to measure; something like your group by or merge...
t1 = time.time() # check clock after
time_to_run_step = t1 - t0

Medindo quanto tempo cada etapa leva para ser executada, você pode concentrar sua energia otimizando as etapas mais lentas. Por exemplo, otimizar uma operação de 0,1 segundo para ser 100x mais rápido é menos bom do que otimizar uma operação de 10 segundos para ser 2x mais rápido.

Ao pensar em como otimizar seu código, duas questões surgiram:

1. Podemos aplicar as prioridades em ordem inversa para evitar a filtragem de prioridades já categorizadas?
2. Podemos realizar todo o grupo trabalhando ao mesmo tempo?

Agrupar por e mesclar são operações caras, pois geralmente são dimensionadas de forma quadrática (nº de faturas X nº de códigos). Aposto que essas são as etapas mais lentas em seu código, mas você deve cronometrar para verificar.

# Act 1: set up everything for the big group by
# priority 1
# will be setup at the end of Act 2

# priority 2
df['milk'] = df['code'].isin(MILK)

# priority 3.a
# priority 3.b
juice = df['code'].isin(JUICE)
df['juice_qty'] = df['qty']
df.loc[~juice, 'juice_qty'] = 0 # I thought df['juice_qty'][~juice] was intuitive, but it gave a warning https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#returning-a-view-versus-a-copy
# distinguish single from many juice in Act 2

# priority 4
df['hot'] = df['code'].isin(HOT)

# priority 5
df['dessert'] = df['code'].isin(DESSERT)


# Act 2: the big group by and merge
invoices = df.groupby(['invoice']).agg({
    'qty': 'sum',
    'milk': 'any',
    'juice_qty': 'sum',
    'hot': 'any',
    'dessert': 'any',
}).rename(columns={
    'qty': 'total', # this is renamed because joining with duplicate names leads to qty_x and qty_y
    'juice_qty': 'juice_total',
})
# priority 1
invoices['mega'] = invoices['total'] >= 10

# priority 3.a
# priority 3.b
invoices['one_juice'] = invoices['juice_total'] == 1
invoices['many_juice'] = invoices['juice_total'] > 1

df = df.merge(invoices, on='invoice', how='left')


# Act 3: apply the categories
# apply the categories in reverse order to overwrite less important with the more important
df['category'] = 'Others'
df.loc[df['dessert_y'], 'category'] = 'DessertLovers'
df.loc[df['hot_y'], 'category'] = 'HotLovers'
df.loc[df['many_juice'], 'category'] = 'ManyJuice'
df.loc[df['one_juice'], 'category'] = 'OneJuice'
df.loc[df['milk_y'], 'category'] = 'Healthy'
df.loc[df['mega'], 'category'] = 'Mega'

df = df[['invoice', 'code', 'qty', 'category']] # get the columns you care about

@Tommy e @ MaartenFabré notaram um bug na categorização de suco único e muitos. Eu editei esta resposta com uma correção.

Edit: Existem algumas respostas para esta questão abrangendo o estouro de pilha também. Abaixo um resumo em 20/09/2020.

• categorização original baseada em prioridade usando pandas / python
• one_group_by https://codereview.stackexchange.com/a/249481/230673
• np_select https://stackoverflow.com/a/63947686/14308614
• np_select_where https://codereview.stackexchange.com/a/249586/230673
• https://codereview.stackexchange.com/a/249486/230673 não foi plotado porque a complexidade do tempo era diferente

O desempenho foi plotado usando o código de https://stackoverflow.com/a/63947686/14308614

Answer 2

Em vez de agrupar pela fatura em cada categoria, eu inverteria a lógica. Agrupe por fatura e, em seguida, classifique essa fatura.

categories = pd.concat(
    classify_invoice(data) for invoice, data in df.groupby("invoice")
)

|    | 0             |
|---:|:--------------|
|  0 | OneJuice      |
|  1 | OneJuice      |
|  2 | Healthy       |
|  3 | Healthy       |
|  4 | Healthy       |
|  5 | Mega          |
|  6 | Mega          |
|  7 | Mega          |
|  8 | ManyJuice     |
|  9 | ManyJuice     |
| 10 | ManyJuice     |
| 11 | HotLovers     |
| 12 | HotLovers     |
| 13 | DessertLovers |
| 14 | DessertLovers |
| 15 | DessertLovers |
| 16 | Others        |

Então, para adicionar isso ao resultado, você pode atribuir.

result = df.assign(category=categories)

Aqui eu usei assign, que retorna um novo DataFrame. Eu faço isso de propósito, para que você possa manter seu DataFrame original intacto. As alterações no local do DataFrame original podem ser uma fonte de erros.

Classificador

Então, precisamos apenas projetar o classificador. Aqui precisamos de uma função que aceite um DataFrame que cubra exatamente 1 fatura e retorne uma série com a categoria, com o mesmo índice da fatura.

Prioridade 1

A prioridade 1 então é fácil:

def classify_invoice(order: pd.DataFrame) -> pd.Series:
    if order["qty"].sum() > 10:
        return pd.Series("Mega", index=order.index)

Prioridade 2

A prioridade 2 também é muito fácil:

    milk_codes = {106, 107, 108}
    if order["code"].isin(milk_codes).any():
        return pd.Series("Healthy", index=order.index)

Observe que renomeei a variável Milkpara milk_codes, uma vez que descreve melhor o que significa, e que a converti em um set, uma vez que essa é a estrutura de dados destinada a verificações de contenção

outras prioridades

def classify_invoice(order: pd.DataFrame) -> pd.Series:
    if order["qty"].sum() > 10:
        return pd.Series("Mega", index=order.index)

    milk_codes = {106, 107, 108}
    if order["code"].isin(milk_codes).any():
        return pd.Series("Healthy", index=order.index)

    juice_codes = {104, 105}
    juices_amount = order.loc[order["code"].isin(juice_codes), "qty"].sum()
    if juices_amount == 1:
        return pd.Series("OneJuice", index=order.index)
    if juices_amount > 1:
        return pd.Series("ManyJuice", index=order.index)

    hot_codes = {103, 109}
    if order["code"].isin(hot_codes).any():
        return pd.Series("HotLovers", index=order.index)

    dessert_codes = {110, 111}
    if order["code"].isin(dessert_codes).any():
        return pd.Series("DessertLovers", index=order.index)

    return pd.Series("Others", index=order.index)

Testando

Como você transferiu a categorização para outra função, pode testar isso isoladamente

---

Variação

def classify_invoice2(order: pd.DataFrame) -> pd.Series:
    if order["qty"].sum() > 10:
        return "Mega"

    milk_codes = {106, 107, 108}
    if order["code"].isin(milk_codes).any():

        return "Healthy"

    juice_codes = {104, 105}
    juices_amount = order.loc[order["code"].isin(juice_codes), "qty"].sum()
    if juices_amount == 1:
        return "OneJuice"
    if juices_amount > 1:
        return "ManyJuice"

    hot_codes = {103, 109}
    if order["code"].isin(hot_codes).any():
        return "HotLovers"

    dessert_codes = {110, 111}
    if order["code"].isin(dessert_codes).any():
        return "DessertLovers"

    return "Others"

df.join(
    df.groupby("invoice")
    .apply(classify_invoice2)
    .rename("category"),
    on = "invoice"
)

Isso é quase tão rápido quanto minha outra solução e um pouco mais simples de seguir.

micro otimizações

Agora os códigos são definidos a cada grupo. Se houver muitas faturas, pode ser mais rápido defini-las fora do método:

milk_codes = {106, 107, 108}
juice_codes = {104, 105}
hot_codes = {103, 109}
dessert_codes = {110, 111}

def classify_invoice3(order: pd.DataFrame) -> pd.Series:
    if order["qty"].sum() > 10:
        return "Mega"

    if order["code"].isin(milk_codes).any():

        return "Healthy"

    juices_amount = order.loc[order["code"].isin(juice_codes), "qty"].sum()
    if juices_amount == 1:
        return "OneJuice"
    if juices_amount > 1:
        return "ManyJuice"

    if order["code"].isin(hot_codes).any():
        return "HotLovers"

    if order["code"].isin(dessert_codes).any():
        return "DessertLovers"
    return "Others"

categórico

Trabalhar com um categórico pode ser mais rápido do que com uma coluna de strings também:

CATEGORIES = {
    0: "Mega",
    1: "Healthy",
    2: "OneJuice",
    3: "ManyJuice",
    4: "HotLovers",
    5: "DessertLovers",
    6: "Others",
}


def classify_invoice4(order: pd.DataFrame) -> pd.Series:
    if order["qty"].sum() > 10:
        return 0

    if order["code"].isin(milk_codes).any():

        return 1

    juices_amount = order.loc[order["code"].isin(juice_codes), "qty"].sum()
    if juices_amount == 1:
        return 2
    if juices_amount > 1:
        return 3

    if order["code"].isin(hot_codes).any():
        return 4

    if order["code"].isin(dessert_codes).any():
        return 5
    return 6

df.join(
    (
        df.groupby("invoice")
        .apply(classify_invoice4)
        .rename("category")
        .astype(pd.Categorical(list(CATEGORIES)))
        .cat.rename_categories(CATEGORIES)
    ),
    on="invoice",
)

No benchmark com os dados de amostra, isso foi um pouco mais lento, mas para conjuntos de dados maiores, isso pode ser mais rápido

---

entorpecido

Você também pode fazer isso em terreno entorpecido:

def classify_invoice_numpy(invoices, quantities, codes):
    SODA = np.array([101, 102])
    HOT = np.array([103, 109])
    JUICE = np.array([104, 105])  # remember spaces after commas
    MILK = np.array([106, 107, 108])
    DESSERT = np.array([110, 111])

    juices = np.isin(codes, JUICE)
    milk = np.isin(codes, MILK)
    hot = np.isin(codes, HOT)
    dessert = np.isin(codes, DESSERT)

    result = -np.ones(len(invoices), dtype=int)

    for invoice in np.unique(invoices):
        index = invoices == invoice

        if quantities[index].sum() >= 10:
            result[index] = 0
            continue

        if milk[index].any():
            result[index] = 1
            continue

        juices_index = index & juices
        if juices_index.any():
            if quantities[juices_index].sum() == 1:
                result[index] = 2
                continue
            else:
                result[index] = 3
                continue

        if hot[index].any():
            result[index] = 4
            continue

        if dessert[index].any():
            result[index] = 5
            continue

    return result

def solution_maarten_numpy(data):
    return data.assign(
        category=pd.Series(
            classify_invoice_numpy(
                data["invoice"].values,
                data["qty"].values,
                data["code"].values,
            ),
            index=data.index,
        ).map(CATEGORIES)
    )

---

avaliação comparativa

Eu fiz alguns benchmarking

dados fictícios:

def dummy_data(
    n: int = 100, lines_per_invoice: int = 3, seed: int = 0
) -> pd.DataFrame:
    random_generator = np.random.default_rng(seed=seed)
    samples = (
        random_generator.normal(loc=lines_per_invoice, scale=2, size=n)
        .round()
        .astype(int)
    )
    samples = np.where(samples > 0, samples, 1)
    invoices = np.repeat(np.arange(n), samples)
    quantities = random_generator.integers(1, 10, size=len(invoices))
    codes = random_generator.choice(np.arange(101, 112), size=len(invoices))
    return pd.DataFrame(
        {"invoice": invoices, "qty": quantities, "code": codes}
    )

compare quando há algo diferente

def compare_results(left, right):
    differences = (left != right).any(axis=1)
    return left[differences].merge(
        right.loc[differences, "category"], left_index=True, right_index=True
    )

benchmark

def benchmark(functions, size=100, lines_per_invoice=3, seed=0):

    data_original = dummy_data(
        n=size, lines_per_invoice=lines_per_invoice, seed=seed
    )
    yield data_original
    benchmark_result = categorise_dawson(data_original)

    for function in functions:
        data = data_original.copy()
        result = function(data)
        try:
            pd.testing.assert_frame_equal(result, benchmark_result)
        except AssertionError:
            print(f"method {function.__name__} differs from the benchmark")
            #             print(result)
            #             print(benchmark_result)
            print(compare_results(benchmark_result, result))
            #             pd.testing.assert_frame_equal(result, benchmark_result)
            continue
        try:
            pd.testing.assert_frame_equal(data, data_original)
        except AssertionError:
            print(f"method {function.__name__} changes the original data")
            continue

        time = timeit.timeit(
            "function(data)",
            globals={"function": function, "data": data},
            number=1,
        )

        yield function.__name__, time

chamando

data_originals = {}
sizes = 10, 100, 1000, 10000
functions = [
    solution_maarten_1,
    solution_maarten_2,
    solution_maarten_3,
    solution_maarten4,
    solution_maarten_numpy,
    categorise_dawson,
    categorise_OP,
]

result_df = pd.DataFrame(index=[function.__name__ for function in functions])
for size in sizes:
    data_original, *results = benchmark(functions=functions, size=size,)
    data_originals[size] = data_original
    result_df[size] = pd.Series(dict(results))

|                        |        10 |       100 |      1000 |     10000 |
|:-----------------------|----------:|----------:|----------:|----------:|
| solution_maarten_1     | 0.0077566 | 0.089533  | 0.838123  | 9.03633   |
| solution_maarten_2     | 0.0085086 | 0.0564532 | 0.521976  | 5.17024   |
| solution_maarten_3     | 0.0051805 | 0.0461194 | 0.545553  | 6.22027   |
| solution_maarten4      | 0.0091025 | 0.0647327 | 0.545063  | 5.88994   |
| solution_maarten_numpy | 0.0013638 | 0.0038171 | 0.0156193 | 0.977562  |
| categorise_dawson      | 0.0342312 | 0.0253829 | 0.0320662 | 0.0790319 |
| categorise_OP          | 0.0480042 | 0.0463131 | 0.0542139 | 0.150899  |

Então meu código começa mais rápido para tamanhos menores, mas muda quase linearmente com o tamanho, enquanto seu código e @dawsons são quase constantes para tamanho

---

código completo

#!/usr/bin/env python
# coding: utf-8

# In[1]:


import numpy as np
import pandas as pd
import timeit


# In[2]:


def dummy_data(
    n: int = 100, lines_per_invoice: int = 3, seed: int = 0
) -> pd.DataFrame:
    random_generator = np.random.default_rng(seed=seed)
    samples = (
        random_generator.normal(loc=lines_per_invoice, scale=2, size=n)
        .round()
        .astype(int)
    )
    samples = np.where(samples > 0, samples, 1)
    invoices = np.repeat(np.arange(n), samples)
    quantities = random_generator.integers(1, 10, size=len(invoices))
    codes = random_generator.choice(np.arange(101, 112), size=len(invoices))
    return pd.DataFrame(
        {"invoice": invoices, "qty": quantities, "code": codes}
    )


# In[3]:


def compare_results(left, right):
    differences = (left != right).any(axis=1)
    return left[differences].merge(
        right.loc[differences, "category"], left_index=True, right_index=True
    )


# In[63]:


Soda = [101, 102]
Hot = [103, 109]
Juice = [104, 105]
Milk = [106, 107, 108]
Dessert = [110, 111]


def categorise_OP(df):
    # Calculating Priority No.1
    L = df.groupby(["invoice"])["qty"].transform("sum") >= 10
    df_Large = df[L]["invoice"].to_frame()
    df_Large["category"] = "Mega"
    df_Large.drop_duplicates(["invoice"], inplace=True)

    # Calculating Priority No.2
    df_1 = df[~L]  # removing Priority No.1 calculated above
    M = df_1["code"].isin(Milk).groupby(df_1["invoice"]).transform("any")
    df_Milk = df_1[M]["invoice"].to_frame()
    df_Milk["category"] = "Healthy"
    df_Milk.drop_duplicates(["invoice"], inplace=True)

    # Calculating Priority No.3

    # 3.a Part -1

    df_2 = df[~L & ~M]  # removing Priority No.1 & 2 calculated above
    J_1 = (df_2["qty"] * df_2["code"].isin(Juice)).groupby(
        df_2["invoice"]
    ).transform("sum") == 1
    df_SM = df_2[J_1]["invoice"].to_frame()
    df_SM["category"] = "OneJuice"
    df_SM.drop_duplicates(["invoice"], inplace=True)

    # 3.b Part -2
    J_2 = (df_2["qty"] * df_2["code"].isin(Juice)).groupby(
        df_2["invoice"]
    ).transform("sum") > 1
    df_MM = df_2[J_2]["invoice"].to_frame()
    df_MM["category"] = "ManyJuice"
    df_MM.drop_duplicates(["invoice"], inplace=True)

    # Calculating Priority No.4
    df_3 = df[
        ~L & ~M & ~J_1 & ~J_2
    ]  # removing Priority No.1, 2 & 3 (a & b) calculated above
    H = df_3["code"].isin(Hot).groupby(df_3["invoice"]).transform("any")
    df_Hot = df_3[H]["invoice"].to_frame()
    df_Hot["category"] = "HotLovers"
    df_Hot.drop_duplicates(["invoice"], inplace=True)

    # Calculating Priority No.5
    df_4 = df[
        ~L & ~M & ~J_1 & ~J_2 & ~H
    ]  # removing Priority No.1, 2, 3 (a & b) and 4 calculated above
    D = df_4["code"].isin(Dessert).groupby(df_4["invoice"]).transform("any")
    df_Dessert = df_4[D]["invoice"].to_frame()
    df_Dessert["category"] = "DessertLovers"
    df_Dessert.drop_duplicates(["invoice"], inplace=True)

    # merge all dfs
    category = pd.concat(
        [df_Large, df_Milk, df_SM, df_MM, df_Hot, df_Dessert],
        axis=0,
        sort=False,
        ignore_index=True,
    )

    # Final merge to the original dataset
    return df.merge(category, on="invoice", how="left").fillna(value="Others")


# In[7]:


SODA = [101, 102]
HOT = [103, 109]
JUICE = [104, 105]  # remember spaces after commas
MILK = [106, 107, 108]
DESSERT = [110, 111]


def categorise_dawson(df):
    df = df.copy()
    df["milk"] = df["code"].isin(MILK)

    # priority 3.a
    juice = df["code"].isin(JUICE)
    invoices_one_juice = df[juice].groupby("invoice")["qty"].sum() == 1
    df["one_juice"] = (
        invoices_one_juice.reindex(df["invoice"]).fillna(False).values
    )
    # priority 3.b
    df["many_juice"] = (
        (~invoices_one_juice).reindex(df["invoice"]).fillna(False).values
    )

    # priority 4
    df["hot"] = df["code"].isin(HOT)

    # priority 5
    df["dessert"] = df["code"].isin(DESSERT)

    # Act 2: the big group by and merge
    invoices = (
        df.groupby(["invoice"])
        .agg(
            {
                "qty": "sum",
                "milk": "any",
                "one_juice": "any",
                "many_juice": "any",
                "hot": "any",
                "dessert": "any",
            }
        )
        .rename(
            columns={
                "qty": "total",  # this is renamed because joining with duplicate names leads to qty_x and qty_y
            }
        )
    )
    # priority 1
    invoices["mega"] = invoices["total"] >= 10

    df = df.merge(invoices, on="invoice", how="left")

    # Act 3: apply the categories
    # apply the categories in reverse order to overwrite less important with the more important
    df["category"] = "Others"
    df.loc[df["dessert_y"], "category"] = "DessertLovers"
    df.loc[df["hot_y"], "category"] = "HotLovers"
    df.loc[df["many_juice_y"], "category"] = "ManyJuice"
    df.loc[df["one_juice_y"], "category"] = "OneJuice"
    df.loc[df["milk_y"], "category"] = "Healthy"
    df.loc[df["mega"], "category"] = "Mega"

    return df[
        ["invoice", "qty", "code", "category"]
    ]  # get the columns you care about


# In[72]:


def classify_invoice1(order: pd.DataFrame) -> pd.Series:
    if order["qty"].sum() >= 10:
        return pd.Series("Mega", index=order.index)

    milk_codes = {106, 107, 108}
    if order["code"].isin(milk_codes).any():
        return pd.Series("Healthy", index=order.index)

    juice_codes = {104, 105}
    juices_amount = order.loc[order["code"].isin(juice_codes), "qty"].sum()

    if juices_amount == 1:
        return pd.Series("OneJuice", index=order.index)
    if juices_amount > 1:
        return pd.Series("ManyJuice", index=order.index)

    hot_codes = {103, 109}
    if order["code"].isin(hot_codes).any():
        return pd.Series("HotLovers", index=order.index)

    dessert_codes = {110, 111}
    if order["code"].isin(dessert_codes).any():
        return pd.Series("DessertLovers", index=order.index)

    return pd.Series("Others", index=order.index)


def solution_maarten_1(data):
    categories = pd.concat(
        classify_invoice1(data) for invoice, data in data.groupby("invoice")
    )
    return data.assign(category=categories)


# In[14]:


def classify_invoice2(order: pd.DataFrame) -> pd.Series:
    if order["qty"].sum() >= 10:
        return "Mega"

    milk_codes = {106, 107, 108}
    if order["code"].isin(milk_codes).any():

        return "Healthy"

    juice_codes = {104, 105}
    juices_amount = order.loc[order["code"].isin(juice_codes), "qty"].sum()
    if juices_amount == 1:
        return "OneJuice"
    if juices_amount > 1:
        return "ManyJuice"

    hot_codes = {103, 109}
    if order["code"].isin(hot_codes).any():
        return "HotLovers"

    dessert_codes = {110, 111}
    if order["code"].isin(dessert_codes).any():
        return "DessertLovers"

    return "Others"


def solution_maarten_2(data):
    return data.join(
        data.groupby("invoice").apply(classify_invoice2).rename("category"),
        on="invoice",
    )


# In[17]:


milk_codes = {106, 107, 108}
juice_codes = {104, 105}
hot_codes = {103, 109}
dessert_codes = {110, 111}


def classify_invoice3(order: pd.DataFrame) -> pd.Series:
    if order["qty"].sum() >= 10:
        return "Mega"

    if order["code"].isin(milk_codes).any():
        return "Healthy"

    juices_amount = order.loc[order["code"].isin(juice_codes), "qty"].sum()
    if juices_amount == 1:
        return "OneJuice"
    if juices_amount > 1:
        return "ManyJuice"

    if order["code"].isin(hot_codes).any():
        return "HotLovers"

    if order["code"].isin(dessert_codes).any():
        return "DessertLovers"
    return "Others"


def solution_maarten_3(data):
    return data.join(
        data.groupby("invoice").apply(classify_invoice3).rename("category"),
        on="invoice",
    )


# In[20]:


CATEGORIES = {
    0: "Mega",
    1: "Healthy",
    2: "OneJuice",
    3: "ManyJuice",
    4: "HotLovers",
    5: "DessertLovers",
    -1: "Others",
}


def classify_invoice4(order: pd.DataFrame) -> pd.Series:
    if order["qty"].sum() >= 10:
        return 0

    if order["code"].isin(milk_codes).any():
        return 1

    juices_amount = order.loc[order["code"].isin(juice_codes), "qty"].sum()
    if juices_amount == 1:
        return 2
    if juices_amount > 1:
        return 3

    if order["code"].isin(hot_codes).any():
        return 4

    if order["code"].isin(dessert_codes).any():
        return 5
    return -1


def solution_maarten4(data):
    return data.join(
        (
            data.groupby("invoice")
            .apply(classify_invoice4)
            .map(CATEGORIES)
            .rename("category")
        ),
        on="invoice",
    )


# In[24]:


def classify_invoice_numpy(invoices, quantities, codes):
    SODA = np.array([101, 102])
    HOT = np.array([103, 109])
    JUICE = np.array([104, 105])  # remember spaces after commas
    MILK = np.array([106, 107, 108])
    DESSERT = np.array([110, 111])

    juices = np.isin(codes, JUICE)
    milk = np.isin(codes, MILK)
    hot = np.isin(codes, HOT)
    dessert = np.isin(codes, DESSERT)

    result = -np.ones(len(invoices), dtype=int)

    for invoice in np.unique(invoices):
        index = invoices == invoice

        if quantities[index].sum() >= 10:
            result[index] = 0
            continue

        if milk[index].any():
            result[index] = 1
            continue

        juices_index = index & juices
        if juices_index.any():
            if quantities[juices_index].sum() == 1:
                result[index] = 2
                continue
            else:
                result[index] = 3
                continue

        if hot[index].any():
            result[index] = 4
            continue

        if dessert[index].any():
            result[index] = 5
            continue

    return result


# In[25]:


def solution_maarten_numpy(data):
    return data.assign(
        category=pd.Series(
            classify_invoice_numpy(
                data["invoice"].values,
                data["qty"].values,
                data["code"].values,
            ),
            index=data.index,
        ).map(CATEGORIES)
    )


# In[28]:


import timeit


# In[52]:


def benchmark(functions, size=100, lines_per_invoice=3, seed=0):

    data_original = dummy_data(
        n=size, lines_per_invoice=lines_per_invoice, seed=seed
    )
    yield data_original
    benchmark_result = categorise_dawson(data_original)

    for function in functions:
        data = data_original.copy()
        result = function(data)
        try:
            pd.testing.assert_frame_equal(result, benchmark_result)
        except AssertionError:
            print(f"method {function.__name__} differs from the benchmark")
            #             print(result)
            #             print(benchmark_result)
            print(compare_results(benchmark_result, result))
            #             pd.testing.assert_frame_equal(result, benchmark_result)
            continue
        try:
            pd.testing.assert_frame_equal(data, data_original)
        except AssertionError:
            print(f"method {function.__name__} changes the original data")
            continue

        time = timeit.timeit(
            "function(data)",
            globals={"function": function, "data": data},
            number=1,
        )

        yield function.__name__, time


# In[89]:


data_originals = {}
sizes = 10, 100, 1000, 10000
functions = [
    solution_maarten_1,
    solution_maarten_2,
    solution_maarten_3,
    solution_maarten4,
    solution_maarten_numpy,
    categorise_dawson,
    categorise_OP,
]

result_df = pd.DataFrame(index=[function.__name__ for function in functions])
for size in sizes:
    data_original, *results = benchmark(functions=functions, size=size,)
    data_originals[size] = data_original
    result_df[size] = pd.Series(dict(results))


# In[94]:


print(result_df.to_markdown())


# In[99]:


result_df.T.plot(logx=True, logy=True)

Answer 3

Aqui, forneço uma abordagem diferente para resolver esse problema com mais eficiência. Em comparação com a solução da OP, a otimização primária vem nos seguintes aspectos:

• Ligar isinquatro vezes para cada classe de item (Sobremesa, Quente, Suco, Leite) é ineficiente. Uma abordagem melhor é joinusar o DataFrame original dfcom um Seriesque mapeia cada item para uma classe e, em seguida, aplica pd.get_dummies- se à nova coluna de classe para realizar a codificação one-hot. Minha solução irá operar nas informações da aula diretamente, portanto a segunda etapa não é necessária.

• Cada classe de item recebe um valor de prioridade que está alinhado com sua prioridade na lógica de cálculo do categoryvalor, ou seja, Sobremesa <Quente <Suco <Leite. A lógica de computação poderia ser reescrita da seguinte forma:

  1. Calcule a quantidade total, a quantidade total de suco e o valor máximo de prioridade de cada fatura;
  2. Se a quantidade total> 10, o valor da categoria é "Mega";
  3. Se o valor máximo de prioridade for "Suco" e a quantidade total> 1, o valor da categoria será "Muitos Suco";
  4. Caso contrário, atribua o valor da categoria com base no valor de prioridade máxima.

  Na implementação, a categorycoluna é do tipo categórica INVOICE_TYPEe cada valor de categoria possui seu código numérico correspondente. O valor de prioridade de cada classe de item é o código numérico da categoria correspondente da classe.

• np.selecté utilizado para implementar a lógica if-elif-else de maneira vetorial. (Observação: para a lógica if-else, np.where/ pd.DataFrame.wherepoderia ser utilizado em seu lugar.)

Solução:

import pandas as pd
import numpy as np


def add_category(df: pd.DataFrame, mega_threshold: int = 10):
    # Invoice categories
    INVOICE_TYPE = pd.CategoricalDtype([
        "Others", "DessertLovers", "HotLovers", "ManyJuice", "OneJuice", "Healthy", "Mega"
    ], ordered=True)
    CODE_OTHERS = 0  # Numerical code of 'Others' category

    # Mapping from item classes to invoice category codes
    class_values = pd.Series(
        pd.Categorical(["DessertLovers", "HotLovers", "OneJuice", "Healthy"], dtype=INVOICE_TYPE).codes,
        index=["Dessert", "Hot", "Juice", "Milk"]
    )

    # Mapping from item codes to class priority values, which are equivalent to corresponding invoice category codes
    item_code_values = pd.Series(
        class_values[["Hot", "Juice", "Juice", "Milk", "Milk", "Milk", "Hot", "Dessert", "Dessert"]].to_numpy(),
        index=pd.RangeIndex(103, 112), name="item_value"
    )

    df_item_values = df.join(item_code_values, on="code")
    df_item_values["juice_qty"] = (df_item_values["item_value"] == class_values["Juice"]) * df_item_values["qty"]

    # Compute total quantity, total juice quantity, and maximum item priority value of each invoice by aggregation
    df_invoice_info = df_item_values.groupby("invoice").agg({
        "qty": "sum",
        "juice_qty": "sum",
        "item_value": "max"
    })
    df_invoice_info.columns = ["total_qty", "total_juice_qty", "max_item_value"]

    ## This version of aggregation has better readability but it turns out to be 2~3 times slower than the above
    # df_invoice_info = df_item_values.groupby("invoice").agg(
    #     total_qty=("qty", "sum"),
    #     total_juice_qty=("juice_qty", "sum"),
    #     max_item_value=("item_value", "max")
    # )

    max_invoice_item_values = df_invoice_info["max_item_value"]
    max_invoice_item_values.fillna(CODE_OTHERS, inplace=True, downcast="int8")
    is_mega = df_invoice_info["total_qty"] > mega_threshold
    is_many_juice = ((max_invoice_item_values == class_values["Juice"]) &
                     (df_invoice_info["total_juice_qty"] > 1))

    # Compute invoice category codes
    invoice_type_codes = pd.Series(np.select(
        [is_mega, is_many_juice],
        pd.Categorical(["Mega", "ManyJuice"], dtype=INVOICE_TYPE).codes,
        max_invoice_item_values),
        index=df_invoice_info.index
    )

    # Join category codes with the original DataFrame and transform them to the categorical type INVOICE_TYPE
    df["category"] = pd.Categorical.from_codes(invoice_type_codes[df["invoice"]], dtype=INVOICE_TYPE)

    # For performance testing, returning a copy of df instead of modifying it in-place
    # return df.assign(category=pd.Categorical.from_codes(invoice_type_codes[df["invoice"]], dtype=INVOICE_TYPE))

if __name__ == "__main__":
    df = pd.DataFrame({
        'invoice': [1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4, 5, 5, 6, 6, 6, 7],
        'code': [101, 104, 105, 101, 106, 106, 104, 101, 104, 105, 111, 109, 111, 110, 101, 114, 112],
        'qty': [2, 1, 1, 3, 2, 4, 7, 1, 1, 1, 1, 4, 2, 1, 2, 2, 1]
    })
    add_category(df)
    print(df)

Resultado:

    invoice  code  qty       category
0         1   101    2       OneJuice
1         1   104    1       OneJuice
2         2   105    1        Healthy
3         2   101    3        Healthy
4         2   106    2        Healthy
5         3   106    4           Mega
6         3   104    7           Mega
7         3   101    1           Mega
8         4   104    1      ManyJuice
9         4   105    1      ManyJuice
10        4   111    1      ManyJuice
11        5   109    4      HotLovers
12        5   111    2      HotLovers
13        6   110    1  DessertLovers
14        6   101    2  DessertLovers
15        6   114    2  DessertLovers
16        7   112    1         Others

Código de teste de desempenho para execução do Jupyter Notebook (na add_categoryfunção, uma cópia de dfé retornada em vez da modificação no local) vs. a solução de @JoshDawson e esta solução no SO :

df = pd.DataFrame({
    'invoice': [1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4, 5, 5, 6, 6, 6, 7],
    'code': [101, 104, 105, 101, 106, 106, 104, 101, 104, 105, 111, 109, 111, 110, 101, 114, 112],
    'qty': [2, 1, 1, 3, 2, 4, 7, 1, 1, 1, 1, 4, 2, 1, 2, 2, 1]
})

# Test input DataFrame from OP
test_input = df

%timeit add_category(test_input)
%timeit add_category_dawson(test_input)
%timeit add_category_SO(test_input)

# Test input constructed by duplicating the original DataFrame 10**5 times
# and modifying the output to differentiate the invoice ids in each copy
test_input = pd.concat([df] * 10**5, ignore_index=True)
test_input["invoice"] += test_input.index // df.shape[0] * df["invoice"].max()

%timeit add_category(test_input)
%timeit add_category_dawson(test_input)
%timeit add_category_SO(test_input)

Resultados do teste de desempenho no DataFrame original do OP:

11.9 ms ± 422 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
17.5 ms ± 357 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
9.52 ms ± 106 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Resultados do teste de desempenho em DataFrame grande:

411 ms ± 3.65 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
1 s ± 5.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
1.1 s ± 10.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)