表記の簡素化のため、缶切り不要な缶を「不要缶」、必要な缶を「必要缶」とする。
缶切りも1つの品物で、3種類の品物があるところが厄介。
選ぶ品物の種類1つの個数を固定して、残りを差分更新などで $O(1)$ などで求められればよい。
「同じ品物同士」の中からは、単純に $X_i$ の大きい方から選んだ方が得。
まず、缶切り0個の場合、不要缶のみを $X_i$ の大きい方から最大 $M$ 個、選べる。
缶切りを1つ追加($X$ 個まで開けれる)すると、缶は合計 $M-1$ 個と選べる数が1減るが、
必要缶を追加で $X$ 個まで{新たに選ぶ or 不要缶と入れ替える}ことができる。
よって、
を用意して、缶切りを1つ増やすたびに
$|S|$ が選べる缶の個数をオーバーしてたら1つ減らし、捨てる
$|S|$ が足りなければ、$T$ から必要缶を「選べる缶の合計個数」も「必要缶を開けられる個数」もオーバーしない範囲で増やせる
$|S|$ =選べる缶の合計個数となったら、「必要缶を開けられる個数」をオーバーしない範囲で、$S$ と $T$ の先頭を見て、$S.top < T.top$ なら入れ替えられる
として、差分を更新していくとよい。
1.に関して、缶切りを増やし続ける限り「選べる缶の合計個数」は減り続けるし「選べる必要缶の個数」は増え続ける。
popした缶は、不要缶でも必要缶でも、もうそれは選ばれることは無い。
特に、$T.top$ がpopした缶以下だった場合は、後から追加できる缶でもう改善することはないため、打ち切っても良い
2.に関して、足りないということは、不要缶は既に全て選ばれているので、そちらから選ぶことは考慮しなくてよい。
3.に関して、$T$ は大きい方から追加していっているので $S$ に入っている不要缶は、必ず $T$ 以上。
よって $S.top < T.top$ の場合、$S.top$ は必ず不要缶。
などの枝刈りにより若干実装は簡潔にできるが、このあたりを考察するよりは実装した方が速い or 確実という場合もある。
Python3
from heapq import heapify, heappush, heappop
n, m = map(int, input().split())
INF = 1 << 60
easy_open = []
hard_open = []
openers = []
cur_score = 0
for _ in range(n):
t, x = map(int, input().split())
if t == 0:
easy_open.append(x)
elif t == 1:
hard_open.append(x)
else:
openers.append(x)
easy_open.sort(reverse=True)
gotten = easy_open[:m]
heapify(gotten)
hard_open.sort()
openers.sort(reverse=True)
opener_available = 0
opener_used = 0
cur_score = sum(gotten)
best = cur_score
for opi, opx in enumerate(openers):
can_m = m - opi - 1
if can_m <= 0:
break
if not hard_open:
break
opener_available += opx
if len(gotten) > can_m:
x = heappop(gotten)
cur_score -= x
if x >= hard_open[-1]:
break
while len(gotten) < can_m and opener_used < opener_available and hard_open:
x = hard_open.pop()
cur_score += x
heappush(gotten, x)
opener_used += 1
while opener_used < opener_available and hard_open:
if gotten[0] < hard_open[-1]:
x = hard_open.pop()
y = heappop(gotten)
cur_score += x
cur_score -= y
heappush(gotten, x)
opener_used += 1
else:
break
best = max(best, cur_score)
print(best)
「単純パスが存在しない」ということは、「3頂点は、それとは別のある1頂点から出る別々の3辺の先にある」ことと同じである。
ⓘ←○←◎→○→ⓙ→○
↓ `→○
○→ⓚ
2頂点 $i,j$ を固定すると、$i,j$ を結ぶ単純パス上の点、および「$i$ の、$j$ から遠ざかる側の部分木」やその逆の点は、
$k$ として選ばれると一直線になってしまう。
それに該当しないのは、「$i,j$ を結ぶ単純パス上(両端除く)から分岐して他の枝に入る点」となる。
逆にそのような $i,j,k$ に対し、中央の点(◎)は必ず1点に決まる。(前述の“分岐点”がそれとなる)
よって、次数 $d \ge 3$ の頂点について、各部分木の先にある頂点数 $w_1,w_2,...,w_d$ から3つを選んで掛け合わせればよい。
全ての3つ組の積の総和が、その頂点を中央とする $i,j,k$ の選び方となる。
で、3つ組の積の総和をどうやって求めるかだが、以下の式変形が利用できる。
総和の3乗から、3乗の和や2乗の和を引いたりすると、「別々の3つ」を選んだ合計が残る。
具体的には、$w_1,w_2,...$ の総和を $W_1$、2乗の和を $W_2$、3乗の和を $W_3$ とすると、
となる。適当な頂点を根とした、各頂点の部分木の頂点数を数える木DPで求められる。
Python3
n = int(input())
links = [set() for _ in range(n)]
for _ in range(n - 1):
u, v = map(int, input().split())
u -= 1
v -= 1
links[u].add(v)
links[v].add(u)
dp = [-1] * n
parents = [-1] * n
q = [0]
ans = 0
while q:
u = q[-1]
if dp[u] == -1:
dp[u] = -2
for v in links[u]:
if dp[v] == -1:
q.append(v)
parents[v] = u
else:
q.pop()
p = parents[u]
if len(links[u]) < 3:
res = 1
for v in links[u]:
if v != p:
res += dp[v]
dp[u] = res
else:
sum1 = 0
sum2 = 0
sum3 = 0
for v in links[u]:
if v != p:
x = dp[v]
sum1 += x
sum2 += x * x
sum3 += x * x * x
dp[u] = sum1 + 1
x = n - sum1 - 1
sum1 += x
sum2 += x * x
sum3 += x * x * x
ans += (sum1 ** 3 - 3 * sum1 * sum2 + 2 * sum3) // 6
print(ans)