Codeforces: D2. Counting Is Fun (Hard Version)

Question

The code works okay for the following problem.

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≤𝑡≤10³) — 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≤𝑘≤𝑛, 10⁸<𝑝<10⁹) — the length of the array 𝑎, the upper bound on the elements of 𝑎 and modulus 𝑝.

It is guaranteed that the sum of 𝑛² over all test cases does not exceed 10⁷, 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 in python, but passes with C++.

• I'm aware that the variable naming is pretty awful. In my defense, this isn't for any software, this is just for competitive programming purposes and vars and functions should be concise and easy to view. Let's review the code for the efficiency of the algo.

Answer 1

In my defense, this isn't for any software,

That's absurd.

Source code is for communicating technical ideas to humans. Including collaborators distributed across the internet. Including your future self a few months from now. Be respectful of their time. Communicate carefully.

Source code is for communicating technical ideas to machines. Sometimes the machine computes a result slowly, giving TLE. Be respectful of the machine's time. Communicate carefully.

duplicate code

You posted the same same code twice twice, with the "running" section adding a Solution. It would suffice to ask ask people to read the code just once.

imports

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

Yes, I'm sure that these functions and brown paper packages tied up with strings are some of your favorite things. But they're not germane to the current code, so elide them. If your IDE lacks such refactoring tools, you may use "$ ruff check --fix *.py".

The impact is you have a time limit to execute this code in contest conditions. Including the imports. Importing code you don't use can take significant elapsed time, with no benefit.

globals

Similarly we see the Moore and Von Neumann neighborhoods defined, but not used. Also a pair of symbols denoting ARIZONA and arizona. Plus some other fun facts unrelated to the contest.

And the lambdas should be defs, if for no other reason than we could give each one a """docstring""".

Pep-8 asked you nicely to spell it line rather than l.

You introduce cond, a symbol with a well-known historic meaning different from what you intend. Recommend you choose a more descriptive name.

numpy

The inner loops ranging over 0 .. k+1 seem like good targets for numba or ndarrays. If contest conditions allow such an import, consider broadcasting operations across the psum and ndp arrays.