Friend groups by matching last k movies

Two customers are friends when the set of their last k movies is the
same. Order inside the window does not matter, and friendship is
transitive by construction (same set as A, same set as B), so the answer is
a partition: group customers by a canonical key for their window.

Brute force. Compare every pair of customers: build both windows as sets
and check equality, then union matching pairs. That is O(n^2 * k) and it
needs union-find or a second pass to turn pairs into groups. It is worth
saying out loud, then dropping, because equality is transitive: there is
no need to compare pairs at all.

Optimal: hash the canonical window.

1. For each customer, take history[-k:]. If the history is shorter than
   k, the customer has no full window and cannot be matched (state this
   assumption; see the edge cases).
2. Turn the window into an order-free key. frozenset(window) is the
   cleanest key in Python. A sorted tuple of the distinct titles works the
   same way in a language without hashable sets.
3. Append the customer to groups[key].
4. Keep only groups with at least two members (a friend group needs a
   friend). Sort members inside each group, then sort the groups, so the
   output is deterministic.

Invariants:

• Two customers share a bucket if and only if their last-k sets are equal,
  so every bucket is exactly one friend group.
• Each customer lands in at most one bucket, so groups are disjoint.

Set semantics mean repeats collapse: [A, A, B] and [A, B, B] both have
the set {A, B}. Confirm this with the interviewer. If they want a
multiset, use tuple(sorted(Counter(window).items())) as the key instead.

from collections import defaultdict, deque


def friend_groups(customer_ids: list[str], movies: list[list[str]], k: int) -> list[list[str]]:
    """Group customers whose last k movies form the same set.

    Customers with fewer than k movies have no full window and are skipped.
    Returns groups of size >= 2, members sorted, groups sorted by first member.
    """
    if k <= 0:
        raise ValueError("k must be positive")
    if len(customer_ids) != len(movies):
        raise ValueError("customer_ids and movies must have the same length")

    buckets: dict[frozenset[str], list[str]] = defaultdict(list)
    for cid, history in zip(customer_ids, movies):
        if len(history) < k:
            continue
        buckets[frozenset(history[-k:])].append(cid)

    groups = [sorted(members) for members in buckets.values() if len(members) > 1]
    groups.sort()
    return groups


class LiveFriendIndex:
    """Follow-up: histories grow over time; keep groups current per event."""

    def __init__(self, k: int) -> None:
        if k <= 0:
            raise ValueError("k must be positive")
        self.k = k
        self.window: dict[str, deque[str]] = {}
        self.key_of: dict[str, frozenset[str]] = {}
        self.buckets: dict[frozenset[str], set[str]] = defaultdict(set)

    def watch(self, cid: str, movie: str) -> None:
        win = self.window.setdefault(cid, deque(maxlen=self.k))
        old = self.key_of.pop(cid, None)
        if old is not None:
            self.buckets[old].discard(cid)
            if not self.buckets[old]:
                del self.buckets[old]
        win.append(movie)  # deque drops the oldest title automatically
        if len(win) == self.k:
            key = frozenset(win)
            self.key_of[cid] = key
            self.buckets[key].add(cid)

    def friends_of(self, cid: str) -> list[str]:
        key = self.key_of.get(cid)
        if key is None:
            return []
        return sorted(self.buckets[key] - {cid})


if __name__ == "__main__":
    ids = ["u1", "u2", "u3", "u4"]
    hist = [
        ["Heat", "Alien", "Up", "Jaws"],
        ["Jaws", "Up", "Alien"],
        ["Alien", "Jaws", "Up"],
        ["Jaws", "Heat", "Up"],
    ]
    assert friend_groups(ids, hist, 3) == [["u1", "u2", "u3"]]
    assert friend_groups(ids, hist, 2) == [["u1", "u3"]]
    assert friend_groups(ids, hist, 1) == [["u3", "u4"]]
    assert friend_groups(ids, hist, 4) == []
    assert friend_groups(["a", "b"], [["X", "X", "Y"], ["X", "Y", "Y"]], 3) == [["a", "b"]]
    assert friend_groups([], [], 2) == []

    live = LiveFriendIndex(2)
    for cid, m in [("a", "X"), ("a", "Y"), ("b", "Y"), ("b", "X")]:
        live.watch(cid, m)
    assert live.friends_of("a") == ["b"]
    live.watch("a", "Z")  # a's window becomes {Y, Z}
    assert live.friends_of("a") == [] and live.friends_of("b") == []
    print("ok")

Let n be the number of customers and k the window size.

• friend_groups: O(n * k) to slice and hash the windows, plus
  O(n log n) to sort members and groups. Space O(n * k) for the keys.
  Slicing reads only the tail of each history, so total history length does
  not matter beyond the input itself.
• LiveFriendIndex.watch: O(k) per event to rebuild and hash the key.
  An incremental hash (XOR or sum of per-title hashes over a multiset
  count) makes it O(1); see the follow-ups.
• friends_of: O(g log g) where g is the group size.
• Space for the live index: O(n * k).

• Example: ids = [u1, u2, u3, u4], histories
  [Heat, Alien, Up, Jaws], [Jaws, Up, Alien], [Alien, Jaws, Up],
  [Jaws, Heat, Up], k = 3 → [[u1, u2, u3]] (u4's set is
  {Jaws, Heat, Up}).
• Same data, k = 2 → [[u1, u3]] (both end with {Up, Jaws}).
• Same data, k = 1 → [[u3, u4]] (both last watched Up).
• Same data, k = 4 → [] (only u1 has four movies, and a group of one
  is not a friend group).
• Repeats collapse: [X, X, Y] and [X, Y, Y], k = 3 → [[a, b]].
• Empty input → [].
• k <= 0 → ValueError.
• Live index, k = 2: a watches X, Y; b watches Y, X → friends. a then
  watches Z → a's set is {Y, Z}, so neither has friends.

1. Return friend pairs instead of groups. Expand each group of size
   g into its g * (g - 1) / 2 pairs. Point out that output can be
   quadratic, so groups are the better API.
2. Partial overlap: friends share at least m of the last k movies.
   This is no longer transitive, so hashing one key fails. Build an
   inverted index movie -> customers who have it in their window, then for
   each movie count co-occurrences per pair in a Counter. Pairs with count
   >= m are friends. Cost is sum over movies of (watchers^2), which is
   bad for blockbusters; cap or skip very popular titles, or use MinHash
   with LSH to find likely pairs, then verify.
3. Order matters (exact last-k sequence). Use tuple(history[-k:]) as
   the key. Everything else is the same.
4. Histories stream in. Keep a deque(maxlen=k) per customer and move
   the customer between buckets on each event (LiveFriendIndex). For
   O(1) updates keep a count map per window and an additive hash of the
   distinct titles (add a title's random 64-bit value when its count goes
   0 → 1, subtract when it goes 1 → 0), and verify the set on collisions.
5. Too many customers for one machine. Map each customer to
   (hash(window_key), customer_id) and shuffle by the hash, so each
   reducer owns whole groups. This is a single MapReduce or Spark
   groupByKey.

Mid-level (L4)

• Asks whether order and repeats matter and what to do with short histories.
• Sees that set equality means grouping by a hashable key, not pairwise
  comparison, and writes working code with a frozenset or sorted tuple.
• Gives correct O(n * k) complexity and tests a few edge cases.

Senior (L5)

• Explains why transitivity makes a hash grouping valid, and why that breaks
  for the "at least m in common" variant.
• Builds the inverted index for the partial-overlap follow-up and names its
  worst case on popular titles.
• Handles streaming updates with a fixed-size window per customer and
  correct bucket moves.

Staff (L6+)

• Proposes incremental hashing for O(1) updates and covers collision
  checks.
• Talks about scale: sharding by window hash, MinHash/LSH for approximate
  overlap, and skipping or down-weighting very popular titles.
• Connects the problem to product use (who counts as a "friend", privacy of
  viewing history, how stale a group may be) and picks an API that fits.