Overview

For each element \(A_i\) of array \(A\), determine whether there exists another element whose absolute difference with \(A_i\) is exactly \(K\), and find the total number of elements that satisfy this condition.

Approach

First, a naive approach that comes to mind is “checking all pairs of students \((i, j)\).” However, this method requires examining approximately \(\frac{N^2}{2}\) pairs, and under the constraint \(N = 2 \times 10^5\), the number of computations reaches tens of billions, which would exceed the time limit (TLE).

Instead, we switch to an approach where, from the perspective of each student \(i\), we directly compute what scores would differ from theirs by \(K\), and then check whether anyone in the entire group has that score. When student \(i\) has score \(A_i\), the target scores to search for are either \(A_i - K\) or \(A_i + K\). For example, if \(A_i = 5, K = 3\), we just need to check whether there exists another student with a score of \(2\) or \(8\).

This search for “whether a specific score exists” can be performed extremely efficiently in \(O(1)\) per query by storing the score data in advance in a set or associative array (hash map, dictionary).

However, the case \(K = 0\) requires special attention. When \(K = 0\), the target score becomes \(A_i - 0 = A_i\), which is the same as the student’s own score. If we simply search for this in a set, we would find “ourselves” and mistakenly conclude that the condition is satisfied. Therefore, when \(K = 0\), we need to check “whether there is another student with the same score,” that is, “whether that score appears 2 or more times in total.”

Algorithm

We handle the cases \(K = 0\) and \(K > 0\) separately.

1. Case \(K = 0\) 1. Prepare an associative array (such as Python’s dict) and count how many times each score appears. 2. For every student \(i\), if the occurrence count of their score \(A_i\) is \(2\) or more, then there exists someone else with the same score, so increment the answer count by \(1\).

2. Case \(K > 0\) 1. Store all elements of array \(A\) in a set (such as Python’s set). 2. For every student \(i\), check whether \(A_i - K\) or \(A_i + K\) exists in the set. If so, increment the answer count by \(1\). (Since \(K > 0\), \(A_i \pm K\) can never equal \(A_i\) itself, so there is no risk of mistakenly finding oneself.)

Complexity

• Time complexity: \(O(N)\)
  • Adding data to the associative array or set takes \(O(N)\), and each search per student takes \(O(1)\), repeated \(N\) times for \(O(N)\). The overall complexity is \(O(N)\), which can be processed well within the time limit.
• Space complexity: \(O(N)\)
  • Storing the distinct scores in an associative array or set uses at most \(O(N)\) memory.

Implementation Notes

• Search speed for data: Using the in operator on a Python list (array) takes \(O(N)\) per query, so it is crucial to use set or dict to speed up searches to \(O(1)\).

• Counting the number of “students” who satisfy the condition: Note that the problem asks for “the total number of students who satisfy the condition, not the number of pairs that satisfy the condition.” In the algorithm above, since we make a judgment with each student \(i\) as the subject, we naturally count the number of students.

  Source Code

import sys

def main():
    input_data = sys.stdin.read().split()
    if not input_data:
        return
    
    N = int(input_data[0])
    K = int(input_data[1])
    A = [int(x) for x in input_data[2:N+2]]
    
    ans = 0
    if K == 0:
        counts = {}
        for a in A:
            counts[a] = counts.get(a, 0) + 1
        for a in A:
            if counts[a] >= 2:
                ans += 1
    else:
        S = set(A)
        for a in A:
            if (a - K) in S or (a + K) in S:
                ans += 1
                
    print(ans)

if __name__ == '__main__':
    main()

---

This editorial was generated by gemini-3.1-pro-thinking.