Two-Stage Ranker

Task

Implement two_stage_ranker(retrieval_scores: list[float], rerank_scores: list[float], candidate_size: int, k: int) -> list[int]. Retrieve a bounded candidate set with first-stage scores, then rank only those candidates with second-stage scores.

Requirements

• retrieval_scores and rerank_scores are equal-length one-dimensional sequences of finite numeric values; position i describes the same item in both sequences.
• candidate_size and k are non-negative integers.
• First retrieve at most min(candidate_size, len(retrieval_scores)) candidates by descending retrieval score.
• Break equal retrieval-score ties by smaller original item index, including ties at the candidate cutoff.
• Then rerank only the retrieved candidates by descending rerank score.
• Break equal rerank-score ties by smaller original item index.
• Return at most min(k, candidate_size, len(retrieval_scores)) original item indices in second-stage rank order.
• An item excluded by retrieval must never appear, even if it has the highest rerank score in the full collection.
• Return an empty list when either candidate_size or k is zero, or when the score sequences are empty.
• If candidate_size exceeds the collection size, retrieve the full collection; if k exceeds the retrieved candidate count, return every candidate.
• Do not mutate either score sequence.

Example

two_stage_ranker(
    [0.95, 0.80, 0.70, 0.10],
    [0.20, 0.90, 0.60, 1.00],
    3,
    2,
)
# [1, 2]