Pick the highest-scoring movie row by top-k similarity

The work happens inside each group. A group's score is the sum of its k highest similarity scores (all of them if the group has fewer than k distinct movies). The answer is the group with the highest score. Ties go to the lowest group index.

Brute force. For each group, look up every movie's score, sort the list in descending order, and sum the first k. With G groups, M total memberships and m the largest group, this costs O(M log m). It is correct and is often the first thing to say aloud.

Better: a size-k min-heap per group. Walk the group and keep a min-heap of the k best scores seen so far. When the heap is full and a new score beats the heap's minimum, replace the minimum. At the end the heap holds the top k, so its sum is the group score. The total cost is O(M log k), which matters when groups are large and k is small (the usual case for a "top few recommendations" score). heapq.nlargest(k, ...) does the same thing, and you can use it once you have shown you know how it works.

Details that change the answer:

• Duplicates inside a group. A movie listed twice must not count twice. De-duplicate with a set before scoring.
• Fewer than k movies. Sum what is there. Do not pad, and do not disqualify the group.
• Negative scores. Similarity can be negative (for example cosine similarity). Top-k still means the k largest values, so a group of one movie with score -0.2 scores -0.2, not 0. The running best must start at negative infinity, not 0.
• Empty groups score 0. If there are no groups at all, return -1.
• Ties are decided by the lower index. Compare with a strict > while scanning left to right.
• Floating point. Sums of floats can differ in the last bit depending on order. If ties matter, sort the top-k before summing, or compare with a small epsilon.

Invariants. After each group is processed, best_index names the lowest-index group with the maximum score seen so far. Inside a group, the heap holds min(k, seen) scores, and they are the largest scores seen so far.

from __future__ import annotations

import heapq
import math
from typing import Sequence


def group_score(similarity: Sequence[float], group: Sequence[int], k: int) -> float:
    """Sum of the k highest similarity scores among the distinct movies in group."""
    if k <= 0:
        return 0.0
    heap: list[float] = []  # min-heap holding the best k scores seen so far
    for movie in set(group):
        if not 0 <= movie < len(similarity):
            raise ValueError(f"unknown movie id {movie}")
        score = similarity[movie]
        if len(heap) < k:
            heapq.heappush(heap, score)
        elif score > heap[0]:
            heapq.heapreplace(heap, score)
    return math.fsum(heap)  # fsum: order-independent, exact rounding


def highest_scoring_group(
    similarity: Sequence[float], groups: Sequence[Sequence[int]], k: int
) -> int:
    """Index of the group with the highest top-k score; ties go to the lowest index.

    Returns -1 when there are no groups.
    """
    best_index = -1
    best_score = -math.inf
    for index, group in enumerate(groups):
        score = group_score(similarity, group, k)
        if score > best_score:  # strict: earlier group wins a tie
            best_index, best_score = index, score
    return best_index


def ranked_groups(
    similarity: Sequence[float], groups: Sequence[Sequence[int]], k: int, top: int
) -> list[tuple[int, float]]:
    """Follow-up: the `top` best groups as (index, score), best first."""
    scored = ((group_score(similarity, g, k), i) for i, g in enumerate(groups))
    best = heapq.nsmallest(top, scored, key=lambda pair: (-pair[0], pair[1]))
    return [(i, s) for s, i in best]


if __name__ == "__main__":
    sim = [0.9, 0.1, 0.4, 0.8, 0.3, 0.75]
    groups = [[0, 1, 2], [3, 5, 4], [1, 4]]
    # Group 0: 0.9 + 0.4 = 1.3; group 1: 0.8 + 0.75 = 1.55; group 2: 0.4.
    assert highest_scoring_group(sim, groups, 2) == 1
    # k = 1: group 0 has 0.9, group 1 has 0.8.
    assert highest_scoring_group(sim, groups, 1) == 0
    # Fewer than k movies: sum everything that is there.
    assert group_score(sim, [1], 3) == 0.1
    # Duplicates count once.
    assert group_score(sim, [0, 0, 0], 3) == 0.9
    # Ties go to the lower index.
    assert highest_scoring_group([0.5, 0.5], [[0], [1]], 1) == 0
    # Negative scores: the best negative group still wins over nothing.
    assert highest_scoring_group([-0.2, -0.7], [[1], [0]], 1) == 1
    # An empty group scores 0 and beats all-negative groups.
    assert highest_scoring_group([-0.2], [[0], []], 1) == 1
    assert highest_scoring_group(sim, [], 2) == -1
    assert ranked_groups(sim, groups, 2, 2) == [(1, 1.55), (0, 1.3)]
    print("ok")

Let G be the number of groups, M the total number of memberships across all groups, and n the number of movies.

• group_score: O(m log k) time for a group of m movies, O(m + k) space (the de-duplication set plus the heap).
• highest_scoring_group: O(M log k) time overall and O(max m + k) extra space, because only one group is held at a time.
• Sort-based brute force: O(M log m) time.
• ranked_groups: O(M log k + G log top) time and O(top) extra space for the result heap.

• similarity = [0.9, 0.1, 0.4, 0.8, 0.3, 0.75], groups = [[0, 1, 2], [3, 5, 4], [1, 4]], k = 2 → 1 (scores 1.3, 1.55, 0.4).
• Same input with k = 1 → 0 (0.9 beats 0.8).
• Same input with k = 10 → 1 (1.85 vs 1.4 vs 0.4; groups smaller than k sum everything).
• Group with repeats [0, 0, 0], k = 3 → score 0.9, not 2.7.
• Exact tie similarity = [0.5, 0.5], groups = [[0], [1]], k = 1 → 0.
• All negative similarity = [-0.2, -0.7], groups = [[1], [0]], k = 1 → 1 (-0.2 > -0.7; a best that starts at 0 would wrongly return -1).
• An empty group among negative groups → the empty group wins with 0.
• groups = [] → -1.
• k = 0 → every group scores 0, so the answer is 0 when at least one group exists.
• A movie ID outside 0..n-1 → ValueError (or say you would skip it; agree with the interviewer first).

1. Return the top T groups, not one. Score every group, then keep a size-T heap keyed on (score, -index). This takes O(M log k + G log T) time. ranked_groups above shows it.
2. Exclude movies the user has already watched, and the source movie itself. Pass a seen set and skip those IDs inside group_score. The heap logic does not change. Ask whether a group whose movies are all excluded should score 0 or be left out.
3. Many queries against the same groups (a different source movie each time, a different k). Groups are fixed and scores change per query, so precomputing group sums does not help. What helps is a reverse index from movie to groups when the similarity vector is sparse: only touch groups that contain a movie with a nonzero score. For a fixed source and many k values, sort each group's scores once and keep prefix sums. Each k query is then O(1) per group.
4. Streaming score updates. If similarity scores change one movie at a time, keep for each group a sorted structure of its scores (a sorted list, or two heaps split at rank k) and its current top-k sum. An update touches only the groups that contain that movie, at O(log m) each. Keep a global max structure, such as a heap with lazy deletion, over group scores.
5. Different aggregation. An average of the top k instead of a sum changes the ranking of groups with fewer than k movies. A weighted sum (rank 1 counts most, as in a carousel) needs the top k in sorted order, so sort the heap contents at the end. Neither change affects the cost.

Mid-level (L4)

• Restates the scoring rule and asks about ties, duplicates and groups smaller than k before coding.
• Gets a correct sort-based solution, then improves it to a heap with prompting.
• Tests the worked example and at least one edge case (empty group or k larger than the group).

Senior (L5)

• Goes straight to the size-k heap and states O(M log k) versus O(M log m) without prompting.
• Handles negative scores correctly: the running best starts at negative infinity, and missing movies are not counted as 0.
• Separates group_score from the selection loop so follow-ups 1 and 2 are small changes.

Staff (L6+)

• Frames it as a candidate-ranking step in a recommender and asks about real sizes: number of groups, group length and query rate.
• Reasons about repeated queries (prefix sums, reverse index for sparse scores) and incremental updates, and says when each pays for itself.
• Raises float tie-breaking, stable ordering for reproducible results, and how the aggregation choice (sum vs. average vs. weighted) changes which groups users see.