Problem

An array ???? of ???? non-negative integers is said to be good if all the elements of ???? can be made equal to 0 using the following operation some (possibly, zero) times:
Select two distinct indices ???? and ???? (1≤????<????≤????) and subtract 1 from all ???????? such that ????≤????≤????
. You are given two positive integers ????, ???? and a prime number ????
. Over all (????+1)???? arrays of length ???? such that 0≤????????≤???? for all 1≤????≤????, count the number of good arrays.

Since the number might be too large, you are only required to find it modulo ????.

Input

Each test contains multiple test cases. The first line contains a single integer ????(1≤????≤103) — the number of test cases. The description of the test cases follows.

The first line of each test case contains three positive integers ????, ???? and ???? (3≤????≤3000, 1≤????≤????, 108<????<109) — the length of the array ????, the upper bound on the elements of ???? and modulus ????.

It is guaranteed that the sum of ????2 over all test cases does not exceed 107, and ???? is prime.

Output

For each test case, on a new line, output the number of good arrays modulo ????.

Ideone

import os, sys, getpass, platform
import math, random, decimal, queue, heapq, bisect, itertools, functools, collections, string, operator, timeit, pprint
from queue import Queue, PriorityQueue, LifoQueue
from itertools import accumulate, chain, combinations, combinations_with_replacement, compress, count, cycle, dropwhile, filterfalse, groupby, islice, permutations, product, repeat, starmap, takewhile, tee, zip_longest
from functools import cmp_to_key, lru_cache, partial, partialmethod, reduce, singledispatch, total_ordering
from random import choice, choices, shuffle, sample, random, randint, randrange, uniform, seed, getstate, setstate, getrandbits
from string import ascii_letters, ascii_lowercase, ascii_uppercase, digits, hexdigits, octdigits, punctuation, printable, whitespace
from heapq import heapify, heappush, heappop, heapreplace, nlargest, nsmallest
from bisect import bisect, bisect_left
from collections import defaultdict, OrderedDict, deque, Counter
from math import gcd, factorial, isqrt, comb, perm, prod

cond = False

 
I = lambda: [int(a) for l in sys.stdin for a in l.strip().split()]
S = lambda: [a for l in sys.stdin for a in l.strip().split()]
IM = lambda: [[int(a) for a in l.split()] for l in sys.stdin]
SM = lambda: [[a for a in l.split()] for l in sys.stdin]
D = lambda k=1: {i: list(map(int, input().split())) for i in range(1, 1 + int(input()) * k)}
DS = lambda: {i: [(int(x[0]), x[1]) for _ in range(int(input()))
                  for x in [input().split()]] for i in range(int(input()))}
 
d8 = ((1, 0), (1, 1), (0, 1), (-1, 1), (-1, 0), (-1, -1), (0, -1), (1, -1))
d4 = ((1, 0), (0, 1), (-1, 0), (0, -1))
az, AZ = ascii_lowercase, ascii_uppercase
mod, inf = 1_000_000_007, float('inf')
CNR = lambda n, r: factorial(n) // (factorial(r) * factorial(n - r))
PNR = lambda n, r: factorial(n) // factorial(r)
 
A = IM()
 
if cond:
    print(A)
 
 
def solution(A):
    def f(arr):
        n, k, p = arr
        dp = [[0] * (k + 1) for _ in range(2)]
        dp[0][0] = 1
        prev = [[0] * (k + 1) for _ in range(2)]
 
        for i in range(1, n + 2):
            ndp = [[0] * (k + 1) for _ in range(2)]
            psum = [0, 0]
            for j in range(2):
                for x in range(k + 1):
                    psum[j] += dp[j][x]
                    psum[j] %= p
            for j in range(2):
                a, b = 0, 0
                for x in reversed(range(k + 1)):
                    ndp[j][x] += psum[j]
                    ndp[j][x] += b
                    ndp[j][x] %= p
                    a += prev[j ^ 1][k - x]
                    a %= p
                    b += a
                    b %= p
            prev = dp
            dp = ndp
 
        return (dp[0][0] - dp[1][0] + p) % p
 
    for i in range(1, len(A)):
        s = A[i]
        print(f(s))
 
 
solution(A)

Running the Code

• Type the username of your PC in the cond var.
• Create a file in the same directory as the python file and name it a.txt and paste the input copied below.

import os, sys, getpass, platform
import math, random, decimal, queue, heapq, bisect, itertools, functools, collections, string, operator, timeit, pprint
from queue import Queue, PriorityQueue, LifoQueue
from itertools import accumulate, chain, combinations, combinations_with_replacement, compress, count, cycle, dropwhile, filterfalse, groupby, islice, permutations, product, repeat, starmap, takewhile, tee, zip_longest
from functools import cmp_to_key, lru_cache, partial, partialmethod, reduce, singledispatch, total_ordering
from random import choice, choices, shuffle, sample, random, randint, randrange, uniform, seed, getstate, setstate, getrandbits
from string import ascii_letters, ascii_lowercase, ascii_uppercase, digits, hexdigits, octdigits, punctuation, printable, whitespace
from heapq import heapify, heappush, heappop, heapreplace, nlargest, nsmallest
from bisect import bisect, bisect_left
from collections import defaultdict, OrderedDict, deque, Counter
from math import gcd, factorial, isqrt, comb, perm, prod

cond = getpass.getuser() == 'ADD_YOUR_PC_USERNAME_HERE'

### Create a file and name it 'a.txt' and paste the input in the file:

if cond:
    sys.stdin = open(os.path.join(os.getcwd(), 'a.txt'), 'r')
 
 
I = lambda: [int(a) for l in sys.stdin for a in l.strip().split()]
S = lambda: [a for l in sys.stdin for a in l.strip().split()]
IM = lambda: [[int(a) for a in l.split()] for l in sys.stdin]
SM = lambda: [[a for a in l.split()] for l in sys.stdin]
D = lambda k=1: {i: list(map(int, input().split())) for i in range(1, 1 + int(input()) * k)}
DS = lambda: {i: [(int(x[0]), x[1]) for _ in range(int(input()))
                  for x in [input().split()]] for i in range(int(input()))}
 
d8 = ((1, 0), (1, 1), (0, 1), (-1, 1), (-1, 0), (-1, -1), (0, -1), (1, -1))
d4 = ((1, 0), (0, 1), (-1, 0), (0, -1))
az, AZ = ascii_lowercase, ascii_uppercase
mod, inf = 1_000_000_007, float('inf')
CNR = lambda n, r: factorial(n) // (factorial(r) * factorial(n - r))
PNR = lambda n, r: factorial(n) // factorial(r)
 
A = IM()
 
if cond:
    print(A)
 
 
def solution(A):
    def f(arr):
        n, k, p = arr
        dp = [[0] * (k + 1) for _ in range(2)]
        dp[0][0] = 1
        prev = [[0] * (k + 1) for _ in range(2)]
 
        for i in range(1, n + 2):
            ndp = [[0] * (k + 1) for _ in range(2)]
            psum = [0, 0]
            for j in range(2):
                for x in range(k + 1):
                    psum[j] += dp[j][x]
                    psum[j] %= p
            for j in range(2):
                a, b = 0, 0
                for x in reversed(range(k + 1)):
                    ndp[j][x] += psum[j]
                    ndp[j][x] += b
                    ndp[j][x] %= p
                    a += prev[j ^ 1][k - x]
                    a %= p
                    b += a
                    b %= p
            prev = dp
            dp = ndp
 
        return (dp[0][0] - dp[1][0] + p) % p
 
    for i in range(1, len(A)):
        s = A[i]
        print(f(s))
 
 
solution(A)

Input

4
3 1 998244853
4 1 998244353
3 2 998244353
343 343 998244353

Output

4
7
10
456615865

Question

Is there any way that we could optimize the code so that it would pass the online judge? Now it gives a TLE.