H - Max × Sum Editorial
by
Nyaan
簡単のため、任意の \(i\) \((1 \leq i \leq N-1)\) に対して \(A_i \leq A_{i+1}\) が成り立つことを仮定します。(そうでない場合は、\(A, B\) の列を適切にソートすることで答えが変わらないまま条件を達成できます。)
\(S\) に含まれる最大の要素を \(r\) \((K \leq r)\) と置きます。この時、\(\lbrace 1, 2, \dots, r-1 \rbrace\) から要素を \(K-1\) 個選んで \(S\) に加えることになりますが、上記の条件よりどれを選んでも \(\max_{i \in S} A_i\) の値は \(A_r\) となるので \(A_i\) は式の値に影響が無くなり、\(\sum_{i \in S} B_i\) を最小化することだけ考えればよいです。よってこの場合、式の値は次のようになります。
\[A_r \times \left(B_r + (B_1, B_2, \dots, B_{r-1}\ のうち小さい方から\ K-1\ 個の和)\right)\]
よって、\(r\) を決め打ちして解くことを考えると、
- \(B_1, B_2, \dots, B_{r-1}\) のうち小さい方から \(K-1\) 個の和
が \(r=K, K+1, \dots, N\) において全て計算できればこの問題が解けることが分かりました。
この値は優先度付きキューを用いることで全て計算することが出来ます。多重集合 \(Q_n\) を次のように定義します。
- \(B_1, B_2, \dots, B_n\) のうち小さい方から \(\min(K-1, n)\) 個の要素が入った多重集合
この時、\(Q_1, Q_2, \dots, Q_N\) は次のように差分更新によって求めることが出来ます。(ここで \(\lbrace \lbrace \rbrace \rbrace\) は多重集合を意味する記号です。)
- \(Q_1 = \lbrace \lbrace B_1 \rbrace \rbrace\)
- \(2 \le n \leq K-1\) の時、\(Q_n = Q_{n-1} + \lbrace\lbrace B_n \rbrace\rbrace\) である。
- \(K \leq n \leq N\) の時、\(Q_n\) は \(Q_{n-1} + \lbrace \lbrace B_{n-1} \rbrace \rbrace\) から最大値を取り除いた集合である。
\(Q_n\) の管理に必要な操作は「要素の push」「最大値の pop」なので、\(Q_n\) を優先度付きキューで管理すればこれらの処理を高速に行うことが出来ます。
また、答えの計算に必要なものは \(Q_{K-1}, Q_K, \dots, Q_N\) それぞれの要素の総和です。これらは \(Q_n\) を計算する際に push, pop した要素の値の和も別途管理することで容易に計算できます。
以上の内容を適切に実装することでこの問題を解くことができます。計算量は \(\mathrm{O}(N \log N)\) で十分高速です。
- 実装例(Python)
import heapq
T = int(input())
for t in range(T):
N, K = map(int, input().split())
A = list(map(int, input().split()))
B = list(map(int, input().split()))
C = sorted([(a, b) for a, b in zip(A, B)])
ans = 10**18
Q = []
bsum = 0
for i in range(N):
a, b = C[i]
if len(Q) == K - 1:
ans = min(ans, a * (bsum + b))
heapq.heappush(Q, -b)
bsum += b
if len(Q) > K - 1:
bsum -= -heapq.heappop(Q)
print(ans)
posted:
last update: