Count Pairs | TCS CodeVita 9 Solution (Zone 1) 2020

Count Pairs from TCS CodeVita Season 9 (Zone 1) sounds like a pair-counting problem, but it is really a "nearest neighbour" question in disguise. For every number you need to know whether some other number sits close enough to it. Checking every pair is too slow for 10^5 elements, and the trick that saves you is one of the most useful ideas in competitive programming: after sorting, the closest value to any element is always right next to it.

The problem also has a trap with duplicates, and the older solution on this page fell into it. More on that below.

Problem

You get an array A of N integers and an integer K. Call an element X happy if at least one other element of the array lies in the range [X - K, X + K], in other words, if some other element differs from X by at most K. Print how many elements are happy.

My reading of "other than X itself": it means a different position in the array, not a different value. So if the value 5 appears twice, each 5 has the other 5 as a neighbour at distance 0, and both are happy. The first sample confirms this reading.

Input format: the first line has N and K. The second line has the N array elements.

Output format: one integer, the number of happy elements.

Constraints: 1 <= N <= 10^5, 0 <= K <= 10^5, 0 <= A[i] <= 10^9. Time limit 1 second.

Example 1
Input
6 3
5 5 7 9 15 2
Output
5

Example 2
Input
3 2
1 3 5
Output
3

In the first example only 15 is unhappy, because the nearest other value, 9, is 6 away. In the second, each neighbour is exactly 2 away, and since the range is inclusive, all three count.

Approach

The direct method compares each element with every other element. That is about N squared over 2 comparisons, roughly 5 x 10^9 for N = 10^5, far over a 1-second limit.

Here is the better way to think about it. For a given X, you do not need every other element. You only need the closest one. If even the closest is more than K away, nothing else can be within K.

1. Sort the array.
2. In a sorted array, the closest other value to a[i] is either a[i-1] or a[i+1]. Anything further left is smaller than a[i-1], and anything further right is larger than a[i+1].
3. So a[i] is happy if a[i] - a[i-1] <= K or a[i+1] - a[i] <= K, checking only the neighbours that exist.
4. Count the happy ones.

Sorting does not change the answer, because happiness depends only on values, not on positions.

Dry run on example 1 with K = 3. Sorted: 2 5 5 7 9 15.

• 2: right gap 3, happy
• 5: left gap 3, happy
• 5: left gap 0, happy
• 7: left gap 2, happy
• 9: left gap 2, happy
• 15: left gap 6, no right neighbour, not happy

Total 5, matching the expected output.

Solution in C++

The earlier C++ code here had two real bugs. It read only N from the first line and never read K, so K held garbage. It also removed duplicates with a set before checking, then looked at a[1] and a[n-2] without checking the size, which crashes when all values are equal or N is 1. Removing duplicates also loses the fact that two equal values make each other happy. The version below fixes all of that and is simpler.

#include <bits/stdc++.h>
using namespace std;

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int n;
    long long k;
    cin >> n >> k;                       // read BOTH n and k
    vector<long long> a(n);
    for (auto &x : a) cin >> x;

    // After sorting, the closest values to a[i] sit right next to it.
    sort(a.begin(), a.end());

    int happy = 0;
    for (int i = 0; i < n; i++) {
        bool leftOk  = (i > 0)     && (a[i] - a[i - 1] <= k);
        bool rightOk = (i < n - 1) && (a[i + 1] - a[i] <= k);
        if (leftOk || rightOk) happy++;
    }
    cout << happy << "\n";
    return 0;
}

Edge cases I ran: N = 1 (input 1 5 / 7) prints 0, because a single element has no other element. 2 0 / 4 4 prints 2, which is the duplicate case the old code missed. 3 0 / 1 2 3 prints 0. I also compared it with a brute-force O(N^2) checker on 500 random small arrays and they agreed every time. A random input with N = 10^5 finished in about 20 ms.

Output

5

And 3 for the second example.

Complexity

Time: O(N log N) for the sort, plus O(N) for the scan.

Space: O(N) for the array. Values fit in int, but I used long long so the subtraction is safe even if limits grow.

Common mistakes

• Using strict less than. The range [X - K, X + K] is inclusive, so the check is <= K. Example 2 fails with < K.
• Removing duplicates. A value that appears twice is always happy. Deduplicating throws that information away.
• Touching a[i-1] or a[i+1] at the edges. Check that the neighbour exists before reading it, especially when N is 1.
• Counting pairs instead of elements. Despite the title, the answer is the number of happy elements, not the number of close pairs.

Related: Number Distancing and Railway Station are from the same CodeVita round.

Tip: try to solve the problem yourself before reading the solution. Use the code as a reference to check your approach, and make sure you understand why it works.

More TCS CodeVita solutions All coding problems MCQ practice