atcoder.jp
30分ぐらい。もう少し早く解きたかった。
まず、10100のおかげで、K個取るとき、K+1個取るとき、K+2個取るとき、・・・で、和がかぶることはなく独立して考えることができる。
N＝10とし、このうち3個だけを考えてみたとすると以下の11個の中から自由に3つ取り、何種類の和ができるかを考える。
0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10
すると、最小は一番左の3つを取った時の3,最大は一番右の3つを取った時の27となり、この間はすべて作ることができるため、
差+1の25パターンが存在することができる。
これを同様に4個、5個と計算していけばよい。
左から取ったとき、右から取ったときはあらかじめ累積和で計算しておけばここの部分の計算はO(1)で計算可能なので十分に間に合う。
def solve(N, K):
M = 10 ** 9 + 7
pre = []
post = []
s = 0
for i in range(N + 1):
s += i
pre.append(s)
s = 0
for i in range(N, -1, -1):
s += i
post.append(s)
post = post[::-1]
ans = 0
for i in range(K, N + 2):
a = post[N + 1 - i] - pre[i - 1] + 1
ans += a
ans %= M
return ans
if __name__ == "__main__":
N, K = tuple(map(int, input().split(" ")))
print(solve(N, K))
atcoder.jp
うーむつらい。解説見てコード起こした。
大きい順に左右に振り分ける、というところまでは思いついた、が
局所解が大局解になりえないことに気が付いて（左右が同じ距離だった時局所解は同じだが、そのあとの状況によって大局解として異なる可能性がある）、
DPでも適応するのかな、で終わってしまった。
このDPの適応の仕方は思いつかなかった。。
自分なりの解説の解釈だとこうなった
dp[0][0]
dp[1][0]
dp[2][0]
dp[0][1]
dp[0][2]
dp[1][1]
dp[1][1]は以下のどちらか大きいほうである。
- dp[0][1] + 2番目に大きい値を左に振り分けたときのスコア
- dp[1][0] + 2番目に大きい値を右に振り分けたときのスコア
この時、3番目に以降に大きい値は、1番目と2番目がどのように割り振られたとしても状況は変わらないため、順々に計算をすることができ、
再帰的に以下のように言うことができる
dp[x][y] = max(
dp[x-1][y] + x+y番目に大きい値を左に振り分けたときのスコア,
dp[x][y-1] + x+y番目に大きい値を右に振り分けたときのスコア
)
以下、後述のコード解説。
以下の部分で、すべてを右に振り分けたとすべてを右に振り分けた場合を先に計算しておく。
for a, i in q:
dp[x+1][y] = dp[x][y] + a * abs(i - x)
x += 1
for a, i in q:
dp[x][y+1] = dp[x][y] + a * abs(N - 1 - y - i)
y += 1
右振り分けは右端からの距離の注意して以下の3つに分解できる。
- N - 1 (=一番右端の座標)
- y (=右へ振り分けた後の右端からの距離)
- i (=元の座標)
最初の二つを用いて座標を特定し、元の座標との差を取れば移動距離が出る。
肝は以下の部分。
a, i = q[x+y-1]
left = a * abs(i - (x - 1))
right = a * abs(N - 1 - (y - 1) - i)
dp[x][y] = max(dp[x-1][y] + left, dp[x][y-1] + right)
ans = max(ans, dp[x][y])
x+y番目に大きい値をとってきた後、これを左に振り分けたときと右に振り分けたときのスコアを考える。
ここで注意しなければならないのは、左に振り分けたときの状況はdp[x-1][y]から考えるため、
左に振り分けたときのその座標はx-1となり、移動距離はabs(i - (x-i))となること。
同様に右に振り分けたときは、先ほどの式のうち、yがy-1になる。
def solve(N, A):
q = [(a, i) for i, a in enumerate(A)]
q.sort(reverse=True)
dp = [[-1 for _ in range(N + 1)] for __ in range(N + 1)]
x = 0
y = 0
dp[x][y] = 0
for a, i in q:
dp[x+1][y] = dp[x][y] + a * abs(i - x)
x += 1
x = 0
for a, i in q:
dp[x][y+1] = dp[x][y] + a * abs(N - 1 - y - i)
y += 1
ans = 0
for x in range(1, N+1):
for y in range(1, N+1):
if x + y > N:
continue
a, i = q[x+y-1]
left = a * abs(i - (x - 1))
right = a * abs(N - 1 - (y - 1) - i)
dp[x][y] = max(dp[x-1][y] + left, dp[x][y-1] + right)
ans = max(ans, dp[x][y])
return ans
if __name__ == "__main__":
N = int(input())
A = list(map(int, input().split(" ")))
print(solve(N, A))
dp[x][y]はx+y <=N という制限が付くため、
N+1 x N+1の行列のうち左上三角の部分だけが埋まることになる。
これらのうち最大が答えになる。
