B - チーム編成 / Team Formation Editorial by admin
GPT 5.2 High概要
青木君(選手 \(T\) )を必ず入れたうえで、残り \(K-1\) 人を「青木君以外」から実力値の高い順に選んだときの、実力値合計を求める問題です。
考察
重要なのは、選び方が実はとても単純である点です。
- チームに必ず入るのは青木君(値は \(A_T\))
- 残りは「青木君を除いた \(N-1\) 人」の中から実力値が大きい順に \(K-1\) 個を取ればよい
つまり答えは
\(A_T + (\text{青木君以外の上位 }K-1\text{ 個の合計})\)
になります。
素朴には、青木君以外の配列を作って降順ソートし、先頭から \(K-1\) 個を足せば求まります。しかしソートは \(O(N\log N)\) かかります。\(N \le 2\times 10^5\) なら間に合うことも多いですが、「上位 \(K-1\) 個だけ欲しい」ので、全体をソートするのは無駄が出ます。
そこで、上位 \(K-1\) 個だけを効率よく取り出す方法(ヒープ)を使って、必要な分だけ取り出すようにします。
例:
\(A=[5,1,9,3],\ T=2,\ K=3\) のとき
- 青木君は選手2なので \(A_T=1\)
- 残りは \([5,9,3]\) の上位2個 → \(9,5\)
合計は \(1+9+5=15\)
アルゴリズム
- 青木君の実力値 \(aoki = A[T-1]\) を取り出す。
- \(K=1\) の場合は青木君だけでチームが完成するので \(aoki\) を出力する。
- 青木君以外の実力値配列
othersを作る(\(A\) から \(A[T-1]\) を除いたもの)。 othersの中から大きい順に \(K-1\) 個を取り出し、その合計を求める。- 実装では
heapq.nlargest(K-1, others)を使う(内部でサイズ \(K-1\) のヒープを維持して上位要素を求める)。
- 実装では
aoki + 上位(K-1)個の合計を出力する。
計算量
- 時間計算量: \(O(N \log K)\)
(nlargestで全要素を見ながらサイズ \(K-1\) のヒープを維持するため) - 空間計算量: \(O(N)\)
(othersを作る分が主。ヒープ(上位要素)自体は \(O(K)\))
実装のポイント
\(T\) は1-indexed なので、配列アクセスは
A[T-1]にする。\(K=1\) の特殊ケースに注意(青木君以外を選ばないので、他の処理をせず即出力)。
others = A[:T-1] + A[T:]のようにして、青木君を確実に除外する(うっかり青木君を二重に数えない)。ソースコード
import sys
import heapq
def main():
it = iter(map(int, sys.stdin.buffer.read().split()))
N = next(it)
K = next(it)
T = next(it)
A = [next(it) for _ in range(N)]
aoki = A[T - 1]
if K == 1:
print(aoki)
return
others = A[:T - 1] + A[T:]
top = heapq.nlargest(K - 1, others)
print(aoki + sum(top))
if __name__ == "__main__":
main()
この解説は gpt-5.2-high によって生成されました。
posted:
last update: