Official

D - チームビルディング / Team Building Editorial by admin

Gemini 3.0 Flash

Overview

Given \(N\) employees, select exactly \(K\) of them to form a team. The goal is to maximize the value obtained by subtracting the “total penalty from incompatible pairs within the team” from the “sum of abilities of the selected employees.”

Analysis

The key to solving this problem is to focus on the small constraints.

  • The number of employees \(N\) is at most 18, which is very small.
  • The number of combinations \(\binom{N}{K}\) for choosing \(K\) people from \(N\) is at most \(\binom{18}{9} = 48,620\) in the worst case of \(N=18, K=9\).

Since the number of combinations is sufficiently small, brute-force enumeration of all combinations is an effective approach.

For each combination, we perform the following calculations: 1. Sum up the abilities \(A_i\) of all \(K\) selected employees. 2. For each incompatible pair \((U_j, V_j)\) where both members are among the selected \(K\) employees, subtract the penalty \(B_j\).

We perform this calculation for all combinations and the maximum value among them is the answer.

Algorithm

  1. Organizing penalty information: Store the penalty \(B\) incurred when employees \(u\) and \(v\) are both selected in a two-dimensional array penalty_matrix[u][v]. This allows checking whether a specific pair is incompatible in \(O(1)\).

  2. Enumerating combinations: Use Python’s itertools.combinations to enumerate all patterns of choosing \(K\) people from \(N\).

  3. Calculating the total score: For each combination, compute the score as follows:

    • For each team member \(u\), add current_score += A[u].
    • Then, check whether there is a penalty between \(u\) and another team member \(v\), and if so, subtract current_score -= penalty_matrix[u][v].
    • Be careful not to count the same pair twice.
  4. Updating the maximum: If the computed score is greater than the current maximum, update the value.

Complexity

  • Time complexity: \(O(N^2 + \binom{N}{K} \cdot K^2)\)

    • Initializing the penalty matrix takes \(O(N^2)\).
    • There are \(\binom{N}{K}\) combinations, and for each combination we add the abilities of \(K\) people and perform approximately \(K^2/2\) penalty checks, giving this overall complexity. When \(N=18\), this amounts to at most a few million operations, which is well within the time limit.
  • Space complexity: \(O(N^2)\)

    • The two-dimensional array for managing penalties has size \(N \times N\).

Implementation Notes

  • Index adjustment: The problem gives employee numbers from \(1\) to \(N\), but handling them as \(0\) to \(N-1\) within the program makes array access smoother.

  • Initial value setting: The initial value of the variable max_score that holds the maximum should be set to a large negative value (e.g., \(-2 \times 10^{18}\)). This is because the total penalty may exceed the total ability, making the answer negative.

  • Preventing double counting: When subtracting the penalty for pair \((u, v)\), by adjusting the range of the nested loop (e.g., starting j from i + 1), you can prevent the mistake of subtracting the same pair’s penalty twice.

    Source Code

import sys
from itertools import combinations

def solve():
    # Read all input at once and split by whitespace for efficient parsing
    input_data = sys.stdin.read().split()
    if not input_data:
        return
    
    # Parse N (number of employees), M (number of bad compatibility pairs), 
    # and K (number of employees to be chosen for the team)
    N = int(input_data[0])
    M = int(input_data[1])
    K = int(input_data[2])
    
    # Parse the capability values A_i for each employee
    # A is 0-indexed, so employee i has capability A[i]
    A = [int(x) for x in input_data[3:3+N]]
    
    # penalty_matrix[u][v] stores the reduction in total capability 
    # if both employee u and employee v are in the same team.
    # Employees are converted from 1-indexed to 0-indexed.
    penalty_matrix = [[0] * N for _ in range(N)]
    idx = 3 + N
    for _ in range(M):
        u = int(input_data[idx]) - 1
        v = int(input_data[idx+1]) - 1
        b = int(input_data[idx+2])
        penalty_matrix[u][v] = b
        penalty_matrix[v][u] = b
        idx += 3
        
    # Initialize max_score with a value smaller than any possible team score.
    # The minimum possible score is roughly 1*10^9 - (18*17/2)*10^9 = -152*10^9.
    max_score = -2 * 10**18
    
    # Iterate over all possible ways to choose exactly K employees out of N.
    # itertools.combinations is an efficient iterator implemented in C.
    for combo in combinations(range(N), K):
        current_score = 0
        
        # Calculate the total capability for the chosen subset of employees.
        # We add each employee's capability and subtract penalties for every pair.
        for i, u in enumerate(combo):
            # Add individual capability
            current_score += A[u]
            
            # Retrieve the row of penalties for the current employee u
            p_u = penalty_matrix[u]
            
            # Check all other employees v in the combo that haven't been processed yet
            # to avoid double-counting pairs (u, v) and (v, u).
            for j in range(i + 1, K):
                v = combo[j]
                # Subtract the penalty if both u and v are in the team
                current_score -= p_u[v]
        
        # Update the maximum total capability found across all combinations
        if current_score > max_score:
            max_score = current_score
            
    # Output the final maximum team capability
    print(max_score)

if __name__ == '__main__':
    solve()

This editorial was generated by gemini-3-flash-preview.

posted:
last update: