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}")
def B(S):
if S[0] == "1" or "1" not in S: return "No"
col = [4, 3, 5, 2, 4, 6, 1, 3, 5, 7]
p_col = set()
for p, c in zip(S, col):
if p == "1": p_col.add(c)
return "Yes" if len(p_col) != max(p_col) - min(p_col) + 1 else "No"test_all(B)
def to_column(size, first_pin_col=0):
if size < 1: return []
ret = [first_pin_col]
i = 1 # 2行目のピンから見ていく
cnt = 1
while True:
edge = ret[0] - i
for j in range(i + 1):
if cnt == size: break
ret.append(edge + 2 * j)
cnt += 1
else:
i += 1
continue
break
return ret
②短くしてみた
from math import ceil
from itertools import chain
def to_column(size, first_pin_col=0):
rows = ceil((2 * size + 0.25)**(1/2) - 0.5)
return list(chain.from_iterable(
[[first_pin_col - r + c for c in range(0, (r+1)*2, 2)] for r in range(rows)]
))[:size]
戯れに描画関数も作ってみました。
描画関数定義：print_pins
# 縮小表示
def print_pins(pins, col):
cols = max(col) - min(col) + 1
if cols % 2 == 0: cols += 1
mn = min(col)
col = [c - mn for c in col]
lines = []
i = 0
r = 0
while i != len(pins):
line = [" "] * cols
for _ in range(r + 1):
line[col[i]] = "●" if pins[i] == "1" else "○"
i += 1
if i == len(pins): break
r += 1
lines.append("".join(line))
print(*lines[::-1], sep="\n")
import random
n = 55
pins = random.choices(["0", "1"], k=n)
print_pins(pins, to_column(n))
行が全部埋まらなくてもOK
import random
n = 100
pins = random.choices(["0", "1"], k=n)
print_pins(pins, to_column(n))
上の回答をto_column関数に置き換え
def B(S):
if S[0] == "1" or "1" not in S: return "No"
p_col = set()
for p, c in zip(S, to_column(len(S))):
if p == "1": p_col.add(c)
return "Yes" if len(p_col) != max(p_col) - min(p_col) + 1 else "No"test_all(B)
from itertools import accumulate
def C(N, M, A):
*acc, = accumulate([0] + A)
B = [sum(A[i] * (i + 1) for i in range(M))]
for i in range(N-M):
B.append(B[-1] + M * A[M + i] - (acc[M + i] - acc[i]))
return max(B)test_all(C)
N, M, A = test_dataC[1]["in"]A = np.array(A)
N, M, A# 縮小表示
import pandas as pd
import jinja2
df = pd.DataFrame(
np.stack(
[np.concatenate([[0]*i, A[i:i+N-M+1]*(i+1), [0]*(M-1-i)]) for i in range(M)],
), columns=range(1, 11),
dtype=int
)
df.index = [
"x1",
"x2",
"x3",
"x4",
]
# 縮小表示
style_str = lambda _: "background: lightgreen; font-weight: bold"
sty = df.style
for i, label in enumerate(df.index, 1):
sty = sty.applymap(style_str, subset=(f"{label}", i))
stysum(A[i] * (i + 1) for i in range(M))
次に$A_2$から選んだ場合の値を計算したいのですが…
# 縮小表示
style_str = lambda _: "background: aqua; font-weight: bold"
sty = df.style
for i, label in enumerate(df.index, 2):
sty = sty.applymap(style_str, subset=(f"{label}", i))
sty
def D(N, M, A):
dp = [[0] * (M + 1) for _ in range(N + 1)]
dp[0][1] = -10**14
A = [0] + A
for i in range(1, N+1):
for j in range(1, M+1):
dp[i][j] = (-10**14 if j > i else
max(dp[i-1][j], dp[i-1][j - 1] + A[i] * j))
# print(np.array(dp))
return dp[-1][-1]test_all(D)
AtCoder Inc.
