Market Basket Analysis Project with Python

"""Calculate frequent item pairs, confidence, and lift without extra packages."""

from __future__ import annotations

from itertools import combinations
from collections import Counter
import pandas as pd


def demo_transactions() -> pd.DataFrame:
    baskets = {
        "T001": ["Bread", "Milk", "Eggs"], "T002": ["Bread", "Butter"],
        "T003": ["Milk", "Eggs"], "T004": ["Bread", "Milk"],
        "T005": ["Coffee", "Milk"], "T006": ["Bread", "Eggs"],
        "T007": ["Bread", "Milk", "Butter"], "T008": ["Coffee", "Cookies"],
        "T009": ["Bread", "Milk", "Eggs"], "T010": ["Milk", "Cookies"],
        "T011": ["Bread", "Butter"], "T012": ["Bread", "Milk", "Eggs"],
    }
    return pd.DataFrame([(tid, item) for tid, items in baskets.items() for item in items], columns=["transaction_id", "item"])


def pair_rules(lines: pd.DataFrame, min_support: float = 0.15) -> pd.DataFrame:
    baskets = lines.groupby("transaction_id")["item"].apply(lambda s: sorted(set(s)))
    n = len(baskets)
    if n == 0:
        raise ValueError("At least one transaction is required")
    item_count = Counter(item for basket in baskets for item in basket)
    pair_count = Counter(pair for basket in baskets for pair in combinations(basket, 2))
    rows = []
    for (left, right), count in pair_count.items():
        support = count / n
        if support < min_support:
            continue
        left_support, right_support = item_count[left] / n, item_count[right] / n
        rows.append({"item_a": left, "item_b": right, "support": support,
                     "confidence_a_to_b": support / left_support,
                     "confidence_b_to_a": support / right_support,
                     "lift": support / (left_support * right_support)})
    return pd.DataFrame(rows).sort_values(["lift", "support"], ascending=False).reset_index(drop=True)


def main() -> None:
    print(pair_rules(demo_transactions()).round(3).to_string(index=False))
    print("Lift above 1 indicates co-occurrence above independence, not causation.")


if __name__ == "__main__":
    main()