Official

B - ペアを見つけよう / Find a Pair Editorial by admin Editorial for printing

Qwen3-Coder-480B

Overview

Given each student’s score, count the number of students who have at least one other student with a score difference of exactly \(K\).

Analysis

For this problem, the following naive approach for each student \(i\) would be too inefficient:

  • For student \(i\), compare the score difference with every other student \(j\) (\(j \neq i\)) one by one → The time complexity becomes \(O(N^2)\), and since \(N\) can be up to \(2 \times 10^5\), this won’t finish in time (TLE).

Therefore, the following optimization is needed:

  • By recording how many times each score appears, we can avoid comparing individual students.
  • In particular, when \(K = 0\), we only need to check whether there is another student with the same score.
  • When \(K > 0\), for a given score \(x\), we just need to check whether a student with score \(x - K\) or \(x + K\) exists.

By grouping scores by their values in this way, we can eliminate unnecessary comparisons.

Algorithm

  1. First, count how many times each score appears using collections.Counter.
  2. When \(K = 0\):
    • If there are 2 or more students with the same score, all students with that score satisfy the condition, so add their count.
  3. When \(K > 0\):
    • For each score \(x\), quickly check using a set whether \(x - K\) or \(x + K\) exists.
    • If it exists, all students with score \(x\) satisfy the condition, so add their count.

Example

Input:

5 2
1 3 3 5 7

Score counts: {1: 1, 3: 2, 5: 1, 7: 1}
Score set: {1, 3, 5, 7}

  • Score 1: \(1+2=3\) exists → condition satisfied (1 person)
  • Score 3: \(3-2=1\) and \(3+2=5\) both exist → condition satisfied (2 people)
  • Score 5: \(5-2=3\) exists → condition satisfied (1 person)
  • Score 7: \(7-2=5\) exists → condition satisfied (1 person)

Total: \(1 + 2 + 1 + 1 = 5\) people

Complexity

  • Time complexity: \(O(N)\)
  • Space complexity: \(O(N)\)

The number of distinct scores is at most \(N\), so all counting and set operations can be performed in \(O(N)\).

Implementation Notes

  • By using Counter, we can easily obtain the occurrence count for each score.

  • By using set, we can determine whether a specific score exists in \(O(1)\).

  • The cases \(K = 0\) and \(K > 0\) need to be handled separately.

    Source Code

from collections import Counter

N, K = map(int, input().split())
A = list(map(int, input().split()))

counter = Counter(A)
ans = 0

if K == 0:
    # 同じ点数を持つ人が2人以上いる場合、その点数を持つ全員が条件を満たす
    for score, count in counter.items():
        if count >= 2:
            ans += count
else:
    # K > 0 の場合、Ai - K または Ai + K が存在するかどうかをチェック
    unique_scores = set(A)
    for score in unique_scores:
        if (score - K) in unique_scores or (score + K) in unique_scores:
            ans += counter[score]

print(ans)

This editorial was generated by qwen3-coder-480b.

posted:
last update: