Sparse Attention Count

Problem

Count the nonzero entries of a local attention matrix: each of the n tokens attends to min(window, n) tokens, so the count is n * min(window, n).

Return 0 when n or window is not positive.

Hint

Sparse attention cuts the quadratic n^2 cost down to n * window.

case 1

in: [10,3]

expected: 30

case 2

in: [4,10]

expected: 16

case 3

in: [0,3]

expected: 0

case 4

in: [5,1]

expected: 5

Starter codePython

def sparse_attention_count(n, window):
    # Your code here
    pass

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