[Python]探索(NumPy) ABC165C

ABC165C

１次方程式の条件なので、NumPy配列を使うことで、計算回数を大きく削減できる。NumPy配列がイテラブルのような役割を果たす。

PyPyでも実行可能で、計算速度を上げることができる。（なぜかメモリ使用量が増える）

Python3 1063 ms 106372 KB
PyPy3 571 ms 155196 KB

サンプルコード

import sys
import itertools
import numpy as np
 
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline

N, M, Q = map(int, readline().split())
 
A = np.array(list(itertools.combinations_with_replacement(range(1, M + 1), N)))
# [[1 1 1]
#  [1 1 2]
#  [1 1 3]
#  [1 1 4]
#  [1 2 2] ...

n = len(A)
score = np.zeros(n, np.int32)
m = map(int, read().split())
for a, b, c, d in zip(m, m, m, m):
    cond = A[:, b - 1] - A[:, a - 1] == c # NumPy配列
    score += d * cond
 
print(score.max())