公式

B - チーム編成 / Team Formation 解説 by admin

GPT 5.2 High

Overview

This problem asks you to find the total skill value when you must include Aoki (player \(T\)) on the team, and then select the remaining \(K-1\) members from “everyone except Aoki” in order of highest skill value.

Analysis

The key insight is that the selection method is actually very straightforward.

  • Aoki must be on the team (his value is \(A_T\))
  • For the rest, simply take the \(K-1\) largest skill values from the remaining \(N-1\) people (excluding Aoki)

In other words, the answer is: \(A_T + (\text{sum of the top }K-1\text{ values excluding Aoki})\)

A naive approach would be to create an array excluding Aoki, sort it in descending order, and sum the first \(K-1\) elements. However, sorting costs \(O(N\log N)\). While this is often fast enough for \(N \le 2\times 10^5\), since we only need the top \(K-1\) elements, sorting the entire array is wasteful.

Instead, we use a method that efficiently extracts only the top \(K-1\) elements (a heap), retrieving just the elements we need.

Example: For \(A=[5,1,9,3],\ T=2,\ K=3\): - Aoki is player 2, so \(A_T=1\) - From the remaining \([5,9,3]\), take the top 2 → \(9,5\) Total is \(1+9+5=15\)

Algorithm

  1. Extract Aoki’s skill value \(aoki = A[T-1]\).
  2. If \(K=1\), the team consists of only Aoki, so output \(aoki\).
  3. Create a skill value array others excluding Aoki (i.e., \(A\) with \(A[T-1]\) removed).
  4. Extract the \(K-1\) largest values from others and compute their sum.
    • In the implementation, use heapq.nlargest(K-1, others) (which internally maintains a heap of size \(K-1\) to find the top elements).
  5. Output aoki + sum of top (K-1) elements.

Complexity

  • Time complexity: \(O(N \log K)\) (Because nlargest scans all elements while maintaining a heap of size \(K-1\))
  • Space complexity: \(O(N)\) (Mainly for creating others. The heap (top elements) itself is \(O(K)\))

Implementation Notes

  • \(T\) is 1-indexed, so array access should use A[T-1].

  • Be careful with the special case \(K=1\) (no one besides Aoki is selected, so output immediately without further processing).

  • Use something like others = A[:T-1] + A[T:] to reliably exclude Aoki (be careful not to accidentally double-count Aoki).

    Source Code

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()

This editorial was generated by gpt-5.2-high.

投稿日時:
最終更新: