Katkat

Rank Items by Dot-Product Similarity to a User Embedding, Excluding Seen Items

YahooMachine Learning Engineer · Technical Screen

Last updated

Rank candidate items for a user by the dot product of the user's embedding with each item's embedding, returning item indices from highest to lowest score with ties broken by index. A follow-up removes items the user has already seen, testing vectorized scoring, sorting and filtering.

A recommendation model represents a user and each candidate item as embedding vectors of the same dimension. The similarity score of an item is the dot product of the user's embedding and the item's embedding; a higher score means a more similar item. Given the user's embedding and the embeddings of $N$ items, return the item indices sorted from the highest score to the lowest.

Follow-up from the same round: the user has already seen some of the items. Leave those items out of the ranked result.

In the interview this was to be written with NumPy. Here the inputs are plain lists.

Function Signature

def rank_items(user: list[int], items: list[list[int]], seen: list[int]) -> list[int]:

user is the user's embedding of dimension d, items[i] is the embedding of item i, and seen lists the indices of the items the user has already seen.

Rules

  • The score of item i is sum(user[j] * items[i][j] for j in range(d)).
  • Return the indices of all items that are not in seen, ordered by score from highest to lowest.
  • Items with equal scores are ordered by index, smallest first.
  • If seen is empty, every item is ranked. If every item has been seen, return an empty list.
  • Embedding components are integers in this version, so scores are exact and ties are well defined.

Constraints

  • 1 <= d <= 128, where d = len(user), and every items[i] has length d.
  • 1 <= len(items) <= 10^4
  • Every component of user and of every items[i] is an integer in [-1000, 1000], so every score lies in [-1.28 * 10^8, 1.28 * 10^8].
  • 0 <= len(seen) <= len(items); the indices in seen are distinct, and each lies in [0, len(items) - 1].

Examples

Example 1

Input: user = [1, 2], items = [[3, 0], [1, 1], [0, 2], [-1, 4]], seen = []
Output: [3, 2, 0, 1]

The scores are 3, 3, 4 and 7 for items 0 to 3. Item 3 ranks first and item 2 second. Items 0 and 1 tie at 3, so the smaller index comes first.

Example 2

Input: user = [1, 2], items = [[3, 0], [1, 1], [0, 2], [-1, 4]], seen = [2, 0]
Output: [3, 1]

Same scores as Example 1, but items 2 and 0 have been seen and are removed.

Example 3

Input: user = [2, -1, 0], items = [[1, 1, 5], [0, 3, 1], [2, 4, -3], [-1, -5, 7]], seen = [3]
Output: [0, 2, 1]

The scores are 1, -3, 0 and 3. Item 3 has the highest score but has been seen, so the result ranks items 0, 2 and 1.