Streaming Sink Window Keep Set

Problem

Streaming attention keeps a small cache: the first sink_tokens positions plus the most recent window_tokens positions, with duplicates removed when the two ranges overlap.

The signature is streaming_cache_keep(seq_len, sink_tokens, window_tokens). Return the kept token indices sorted ascending.

Hint

Use a set so the sink and window ranges merge cleanly when they overlap.

Test cases

case 1

in: [10,2,3]

expected: [0,1,7,8,9]

case 2

in: [5,2,10]

expected: [0,1,2,3,4]

case 3

in: [4,0,2]

expected: [2,3]

case 4

in: [3,3,0]

expected: [0,1,2]

case 5

in: [6,4,4]

expected: [0,1,2,3,4,5]

Starter codePython

def streaming_cache_keep(seq_len, sink_tokens, window_tokens):
    # Your code here
    pass

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