NTT Butterfly Step

Problem

Apply one number-theoretic transform butterfly to inputs a and b with twiddle factor w modulo mod.

Return [u + v mod mod, u - v mod mod] where u = a mod mod and v = (w * b) mod mod, with both results normalized to [0, mod). The modulus is assumed positive.

Hint

Reduce the product first, then add and subtract before the final modulus.

Test cases

case 1

in: [1,2,3,7]

expected: [0,2]

case 2

in: [5,0,123,97]

expected: [5,5]

case 3

in: [10,10,1,17]

expected: [3,0]

case 4

in: [0,0,5,11]

expected: [0,0]

case 5

in: [3,4,5,13]

expected: [10,9]

Starter codePython

def ntt_butterfly(a, b, w, mod):
    # Your code here
    pass

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