Beam Search Pruning Count

Problem

Count how many candidate continuations beam search prunes. Each alive beam expands over the whole vocabulary, giving min(beam_width, num_alive) * vocab_size candidates, of which only beam_width survive.

The signature is beam_search_pruning(beam_width, vocab_size, num_alive). Return the number of pruned candidates.

Hint

Pruning keeps the search tractable at the cost of possibly missing the global best sequence.

Test cases

case 1

in: [2,100,2]

expected: 198

case 2

in: [5,1000,5]

expected: 4995

case 3

in: [4,50,3]

expected: 146

case 4

in: [1,10,1]

expected: 9

Starter codePython

def beam_search_pruning(beam_width, vocab_size, num_alive):
    # Your code here
    pass

4 test cases · run your code against them in the editor