Условие

На вход — список уже токенизированных документов: docs = [[token, ...], ...]. Посчитайте idf для каждого уникального токена. Формула: idf(t) = log(N / df(t)), где N — число документов, df(t) — число документов, в которых встречается токен.

Решение

Подход

Идём по документам один раз. Для каждого документа берём set токенов и инкрементируем df[token]. Затем для каждого токена вычисляем log(N/df). Сразу обсуждаем сглаживание log((N+1)/(df+1)) + 1, как в sklearn, чтобы избежать log(0) и не давать бесконечный вес токенам, встречающимся во всех документах.

Реализация

from collections import Counter
from math import log
 
def compute_idf(docs: list[list[str]]) -> dict[str, float]:
    N = len(docs)
    df = Counter()
    for doc in docs:
        for token in set(doc):  # уникальные токены документа
            df[token] += 1
    return {t: log(N / c) for t, c in df.items()}
 
def compute_idf_smooth(docs):
    N = len(docs)
    df = Counter()
    for doc in docs:
        for t in set(doc):
            df[t] += 1
    return {t: log((1 + N) / (1 + c)) + 1 for t, c in df.items()}

Подводные камни

Считать df без set(doc) — типичная ошибка, тогда токен с 5 повторами в одном документе даст df=5 вместо df=1.
Без сглаживания токен, встречающийся во всех документах, даёт idf=0 — что корректно, но если потом домножать на TF в логарифме, можно случайно занулить весь вес.
На больших корпусах хранить df для миллионов токенов — это много памяти; используйте Counter поверх потокового чтения.

Эталонный ответ

df[t] через set каждого документа, затем idf = log(N / df). На проде использовать сглаживание log((1+N)/(1+df)) + 1.

alexeygrigorev/data-science-interviews: посчитать IDF по корпусу