Three NumPy versions of @kcsquared's solution that easily solve the worst case in well under the 10 seconds time limit:
def numpy1(n, k, m):
dp = np.zeros((k+1, m), np.int32)
dp[0][0] = 1
for i in range(1, n+1):
dp[1:,] += dp[:-1, (np.arange(m) + i) % m]
dp %= 10**9 + 7
return dp[k][0]
def numpy2(n, k, m):
dp = np.zeros((k+1, m), np.int32)
dp[0][0] = 1
i = range(m)
for _ in range(n):
i = np.roll(i, 1)
dp[1:,] += dp[:-1, i]
dp %= 10**9 + 7
return dp[k][0]
def numpy3(n, k, m):
dp = np.zeros((k+1, m), np.int32)
dp[0][0] = 1
for i in range(n):
dp[1:,] += np.roll(dp[:-1,], i, axis=1)
dp %= 10**9 + 7
return dp[k][0]
Benchmarks for small, medium and worst case:
n = 19 k = 11 m = 13
-------------------------------------------------------
22856.8 μs 23409.3 μs 23421.4 μs naive
496.9 μs 500.2 μs 524.7 μs dtjc
918.6 μs 928.6 μs 936.3 μs kcsquared
173.5 μs 183.6 μs 191.9 μs numpy1
402.2 μs 403.5 μs 411.1 μs numpy2
297.8 μs 318.4 μs 320.1 μs numpy3
n = 200 k = 100 m = 200
-------------------------------------------------------
2033.6 ms 2177.3 ms 2178.1 ms dtjc
1410.6 ms 1420.2 ms 1430.5 ms kcsquared
19.5 ms 19.8 ms 20.3 ms numpy1
22.5 ms 22.9 ms 23.0 ms numpy2
26.8 ms 27.3 ms 27.3 ms numpy3
n = 1000 k = 100 m = 1000
-------------------------------------------------------
508.0 ms 516.1 ms 519.2 ms numpy1
518.3 ms 518.8 ms 526.3 ms numpy2
495.1 ms 496.4 ms 499.2 ms numpy3
Benchmark code (Try it online!):
from timeit import repeat
from itertools import combinations
from functools import lru_cache
import numpy as np
def naive(n, k, m):
return sum(sum(combi) % m == 0
for combi in combinations(range(1, n+1), k)) % (10**9 + 7)
@lru_cache(None)
def dtjc(n, k, m, r=0):
if k > n:
return 0
if k == 0:
return 1 if r == 0 else 0
return (dtjc(n-1, k, m, r) + dtjc(n-1, k-1, m, (r+n) % m)) % (10**9 + 7)
def kcsquared(max_part_size: int, total_num_parts: int, mod: int) -> int:
BIG_MOD = 10 ** 9 + 7
# Optimization if no partitions exist
if total_num_parts > max_part_size:
return 0
largest_sum = ((max_part_size * (max_part_size + 1) // 2)
- ((max_part_size - total_num_parts) *
(max_part_size - total_num_parts + 1) // 2))
# Optimization if largest partition sum still smaller than mod
if largest_sum < mod:
return 0
dp = [[0 for _ in range(mod)] for _ in range(total_num_parts + 1)]
dp[0][0] = 1
for curr_max_part in range(1, max_part_size + 1):
for curr_num_parts in reversed(range(0, total_num_parts)):
for rem in range(mod):
dp[curr_num_parts + 1][(rem + curr_max_part) % mod] += dp[curr_num_parts][rem]
dp[curr_num_parts + 1][(rem + curr_max_part) % mod] %= BIG_MOD
return dp[total_num_parts][0]
def numpy1(n, k, m):
dp = np.zeros((k+1, m), np.int32)
dp[0][0] = 1
for i in range(1, n+1):
dp[1:,] += dp[:-1, (np.arange(m) + i) % m]
dp %= 10**9 + 7
return dp[k][0]
def numpy2(n, k, m):
dp = np.zeros((k+1, m), np.int32)
dp[0][0] = 1
i = range(m)
for _ in range(n):
i = np.roll(i, 1)
dp[1:,] += dp[:-1, i]
dp %= 10**9 + 7
return dp[k][0]
def numpy3(n, k, m):
dp = np.zeros((k+1, m), np.int32)
dp[0][0] = 1
for i in range(n):
dp[1:,] += np.roll(dp[:-1,], i, axis=1)
dp %= 10**9 + 7
return dp[k][0]
def test(args, solutions, number, format_time):
print('n = %d k = %d m = %d' % args)
print('-' * 55)
for _ in range(1):
results = set()
for func in solutions:
times = sorted(repeat(lambda: dtjc.cache_clear() or results.add(func(*args)), number=number))[:3]
print(*(format_time(t / number) for t in times), func.__name__)
print('results set:', results)
assert len(results) == 1
print()
test((19, 11, 13),
[naive, dtjc, kcsquared, numpy1, numpy2, numpy3],
10,
lambda time: '%7.1f μs ' % (time * 1e6))
test((200, 100, 200),
[dtjc, kcsquared, numpy1, numpy2, numpy3],
1,
lambda time: '%6.1f ms ' % (time * 1e3))
test((1000, 100, 1000),
[numpy1, numpy2, numpy3],
1,
lambda time: '%5.1f ms ' % (time * 1e3))