The Three Little Pigs (Codeforces 1548C)
Tags: Math, fft

Three little pigs from all over the world are meeting for a convention! Every minute, a triple of 3 new pigs arrives on the convention floor. After the $$$n$$$-th minute, the convention ends.The big bad wolf has learned about this convention, and he has an attack plan. At some minute in the convention, he will arrive and eat exactly $$$x$$$ pigs. Then he will get away.The wolf wants Gregor to help him figure out the number of possible attack plans that involve eating exactly $$$x$$$ pigs for various values of $$$x$$$ ($$$1 \le x \le 3n$$$). Two attack plans are considered different, if they occur at different times or if the sets of little pigs to eat are different.Note that all queries are independent, that is, the wolf does not eat the little pigs, he only makes plans!

Input: The first line of input contains two integers $$$n$$$ and $$$q$$$ ($$$1 \le n \le 10^6$$$, $$$1 \le q \le 2\cdot 10^5$$$), the number of minutes the convention lasts and the number of queries the wolf asks.Each of the next $$$q$$$ lines contains a single integer $$$x_i$$$ ($$$1 \le x_i \le 3n$$$), the number of pigs the wolf will eat in the $$$i$$$-th query.

Output: You should print $$$q$$$ lines, with line $$$i$$$ representing the number of attack plans if the wolf wants to eat $$$x_i$$$ pigs. Since each query answer can be large, output each answer modulo $$$10^9+7$$$.

Note: In the example test, $$$n=2$$$. Thus, there are $$$3$$$ pigs at minute $$$1$$$, and $$$6$$$ pigs at minute $$$2$$$. There are three queries: $$$x=1$$$, $$$x=5$$$, and $$$x=6$$$.If the wolf wants to eat $$$1$$$ pig, he can do so in $$$3+6=9$$$ possible attack plans, depending on whether he arrives at minute $$$1$$$ or $$$2$$$.If the wolf wants to eat $$$5$$$ pigs, the wolf cannot arrive at minute $$$1$$$, since there aren't enough pigs at that time. Therefore, the wolf has to arrive at minute $$$2$$$, and there are $$$6$$$ possible attack plans.If the wolf wants to eat $$$6$$$ pigs, his only plan is to arrive at the end of the convention and devour everybody.Remember to output your answers modulo $$$10^9+7$$$!

Sample Inputs:
2 3
1
5
6

5 4
2
4
6
8

Sample Outputs:
9
6
1

225
2001
6014
6939

Solution (Python):
import sys
import __pypy__
int_add = __pypy__.intop.int_add
int_sub = __pypy__.intop.int_sub
int_mul = __pypy__.intop.int_mul

def make_mod_mul(mod=10 ** 9 + 7):
    fmod_inv = 1.0 / mod

    def mod_mul(a, b, c=0):
        res = int_sub(int_add(int_mul(a, b), c), int_mul(mod, int(fmod_inv * a * b + fmod_inv * c)))
        if res >= mod:
            return res - mod
        elif res < 0:
            return res + mod
        else:
            return res
    return mod_mul
mod_mul = make_mod_mul()
mod = 10 ** 9 + 7
stream = sys.stdin
read = stream.readline
(n, q) = map(int, read().split())
sz = 3 * (n + 1)
fac = [0] * (sz + 1)
inv = [0] * (sz + 1)
fac[0] = 1
for i in range(1, sz + 1):
    fac[i] = mod_mul(fac[i - 1], i)
inv[sz] = pow(fac[sz], mod - 2, mod)
for i in range(sz - 1, -1, -1):
    inv[i] = mod_mul(inv[i + 1], i + 1)

def nCr(n, r):
    num = fac[n]
    den = mod_mul(inv[n - r], inv[r])
    return mod_mul(num, den)
coef = [nCr(3 * n + 3, i) for i in range(3 * n + 4)]
coef[0] -= 1
ret = [0] * (sz + 1)
for i in range(3 * n + 3, 2, -1):
    mul = coef[i]
    mul3 = mod_mul(3, mul)
    ret[i - 3] = mul
    coef[i] -= mul
    coef[i - 1] -= mul3
    coef[i - 2] -= mul3
    if coef[i - 1] < 0:
        coef[i - 1] += mod
    if coef[i - 2] < 0:
        coef[i - 2] += mod
outp = [0] * q
for i in range(q):
    x = int(read())
    outp[i] = str(ret[x])
print('\n'.join(outp))

Solution (Python):
import sys, os, io
input = io.BytesIO(os.read(0, os.fstat(0).st_size)).readline

def comb(n, r):
    return fact[n] * inv[r] % mod * inv[n - r] % mod if n >= r >= 0 else 0
(n, q) = map(int, input().split())
mod = pow(10, 9) + 7
l = 3 * n + 5
fact = [1] * (l + 1)
for i in range(1, l + 1):
    fact[i] = i * fact[i - 1] % mod
inv = [1] * (l + 1)
inv[l] = pow(fact[l], mod - 2, mod)
for i in range(l - 1, -1, -1):
    inv[i] = (i + 1) * inv[i + 1] % mod
dp = [n, n, n]
ans0 = [0] * (3 * n + 1)
inv3 = inv[3] * fact[2] % mod
for i in range(1, 3 * n + 1):
    dp0 = [0] * 3
    x = comb(3 * n, i + 1)
    dp0[0] = (x - 2 * dp[0] - dp[1]) % mod * inv3 % mod
    dp0[1] = (x + dp[0] - dp[1]) % mod * inv3 % mod
    dp0[2] = (x + dp[0] + 2 * dp[1]) % mod * inv3 % mod
    dp = dp0
    ans0[i] = (dp[0] + comb(3 * n, i)) % mod
ans = [0] * q
for i in range(q):
    x = int(input())
    ans[i] = ans0[x]
sys.stdout.write('\n'.join(map(str, ans)))

Solution (Python):
import sys, os, io
input = io.BytesIO(os.read(0, os.fstat(0).st_size)).readline

def comb(n, r):
    return fact[n] * inv[r] % mod * inv[n - r] % mod if n >= r >= 0 else 0
(n, q) = map(int, input().split())
mod = pow(10, 9) + 7
l = 3 * n + 5
fact = [1] * (l + 1)
for i in range(1, l + 1):
    fact[i] = i * fact[i - 1] % mod
inv = [1] * (l + 1)
inv[l] = pow(fact[l], mod - 2, mod)
for i in range(l - 1, -1, -1):
    inv[i] = (i + 1) * inv[i + 1] % mod
dp = [n, n, n]
ans0 = [0] * (3 * n + 1)
inv3 = inv[3] * fact[2] % mod
for i in range(1, 3 * n + 1):
    dp0 = [0] * 3
    x = comb(3 * n, i + 1)
    dp0[0] = (x - 2 * dp[0] - dp[1]) % mod * inv3 % mod
    dp0[1] = (x + dp[0] - dp[1]) % mod * inv3 % mod
    dp0[2] = (x + dp[0] + 2 * dp[1]) % mod * inv3 % mod
    dp = dp0
    ans0[i] = (dp[0] + comb(3 * n, i)) % mod
ans = [0] * q
for i in range(q):
    x = int(input())
    ans[i] = ans0[x]
sys.stdout.write('\n'.join(map(str, ans)))