def test_all(f):
for i, data in enumerate(eval(f'test_data{f.__name__[0]}')):
exp = data["out"]
ans = f(*data["in"])
result = "AC" if exp == ans else "WA"
print(f"{i+1} {result}: expected: {exp}, output: {ans}")
import re
def B(S):
print("Yes" if re.match(r'^[A-Z][1-9]\d{5}[A-Z]$', S) else "No")B("Q142857Z") # Yes
B("AB912278C") # No
B('X900000') # No
B('K012345K') # No
def D(N, K, D, a):
dp = [[[-1] * D for _ in range(K + 1)] for _ in range(N + 1)]
for i in range(N + 1):
dp[i][0][0] = 0
# print(np.array(dp))
for i in range(1, N + 1):
for j in range(1, min(i + 1, K + 1)):
for k in range(D):
dp[i][j][k] = max(dp[i][j][k], dp[i - 1][j][k])
x = dp[i - 1][j - 1][k]
if x != -1:
y = x + a[i - 1]
dp[i][j][y % D] = max(y, dp[i - 1][j][y % D])
# print(np.array(dp))
return dp[-1][-1][0]test_all(D)
class BIT:
def __init__(self, a):
if hasattr(a, "__iter__"):
self.size = len(a)
self.data = [0] * self.size
self.bit = [0] * (self.size + 1)
for i, j in enumerate(a):
self.add(i, j)
elif isinstance(a, int):
self.size = a
self.data = [0] * a
self.bit = [0] * (a + 1)
else:
raise TypeError()
def add(self, i, x):
self.data[i] += x
i += 1
while i <= self.size:
self.bit[i] += x
i += i & -i
def sum(self, i):
'''Get sum of data[:i]
'''
total = 0
while i > 0:
total += self.bit[i]
i -= i & -i
return total
def get(self, x):
res = 0
i = 1
while i < (self.size + 1) / 2: i <<= 1
while i > 0:
if res + i < self.size and self.bit[res + i] < x:
x -= self.bit[res + i]
res += i
i >>= 1
res += 1
return res if res >= x else -1
def __getitem__(self, key):
return self.data[key]
def __setitem__(self, key, value):
self.add(key, value - self[key])
self.data[key] = value
def __len__(self):
return self.size
def __repr__(self):
return self.data.__repr__()
def __iter__(self):
return self.data.__iter__()
メイン処理
from operator import itemgetter as ig
def E(N, M, K, A):
# data -> [[sorted index, A index, value]]
data = sorted(([si, i, a] for si, (i, a) in enumerate(sorted(enumerate(A), key=ig(1)))), key=ig(1))
idx = BIT(N)
val = BIT(N)
ans = []
for d in data[:M]:
idx[d[0]] = 1
val[d[0]] = d[2]
ans.append(val.sum(idx.get(K)))
for i, d in enumerate(data[M:]):
idx[d[0]] = 1
val[d[0]] = d[2]
idx[data[i][0]] = 0
val[data[i][0]] = 0
ans.append(val.sum(idx.get(K)))
return " ".join(map(str, ans))test_all(E)
AtCoder Inc.
