Sum of Three Values

Sum of Three Values
CSES medium

You are given an array of \(n\) integers, and your task is to find three values at distinct positions whose sum is \(x\).

Input:

  • First line: two integers \(n\) and \(x\) — the array size and the target sum.
  • Second line: \(n\) integers \(a_1,a_2,\dots,a_n\) — the array values.

Output:

Print three integers: the positions (1-indexed) of the chosen values. If there are several valid triplets, you may print any of them. If there is no such triplet, print

IMPOSSIBLE

Constraints:

  • \(1 \le n \le 5000\)
  • \(1 \le x, a_i \le 10^9\)

Example:

Input:

4 8
2 7 5 1

Output:

1 3 4

Solution

We want three distinct indices \(i,j,k\) such that a_i + a_j + a_k = x.

A fully brute-force solution would check all triples in \(O(n^3)\), which is too slow for \(n = 5000\).

Instead, we use the classic idea:

Brute-force two elements and binary search for the third.


1. Sort while keeping original positions

We first store each element together with its original index, then sort by value:

  • Build an array of pairs: \((a_i, i)\).
  • Sort by \(a_i\).

This lets us use binary search / two-pointer style tricks on the values, while still being able to print the original indices.


2. Fix two indices, binary search the third

After sorting:

  • Loop over all pairs \((i,j)\) with \(0 \le i < j < n\):

    • Let \(\text{need} = x - a_i - a_j\).
    • Now we just need to see if there exists \(k > j\) with \(a_k = \text{need}\).
    • Since the array is sorted, we can binary search for need in the range \((j+1, \dots, n-1)\).

If the binary search finds such a k:

  • We output the original positions of \(a_i, a_j, a_k\).

If we finish checking all pairs with no success, we print IMPOSSIBLE.

Because we always search only in indices strictly greater than \(j\), all three positions are automatically distinct.


3. Complexity

  • Sorting: \(O(n \log n)\)
  • Double loop over \((i, j)\): \(O(n^2)\) pairs
  • For each pair, one binary search: \(O(\log n)\)

Total time:

O(n^2 \log n)

which is fine for \(n \le 5000\).

Memory usage is \(O(n)\) to store the array of pairs.

NoteAlgorithm sketch
  1. Read \(n\), \(x\), and the array.

  2. Build vector<(value, original_index)> and sort by value.

  3. For each i from 0 to n-1:

    • For each j from i+1 to n-1:

      • Let need = x - value[i] - value[j].
      • Binary search need in the range [j+1, n-1].
      • If found at k, print the three original indices and exit.
  4. If nothing found, print IMPOSSIBLE.

#include <iostream>
#include <vector>
using namespace std;

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

    int n;
    long long x;
    cin >> n >> x;

    vector<pair<long long,int>> a(n);
    for (int i = 0; i < n; ++i) {
        long long v;
        cin >> v;
        a[i] = {v, i + 1}; // store value and original 1-based index
    }

    sort(a.begin(), a.end()); // sort by value

    // Fix two indices (i, j) and binary search for the third
    for (int i = 0; i < n; ++i) {
        for (int j = i + 1; j < n; ++j) {
            long long need = x - a[i].first - a[j].first;

            int lo = j + 1, hi = n - 1;
            while (lo <= hi) {
                int mid = (lo + hi) / 2;
                if (a[mid].first == need) {
                    // Found a valid triple
                    cout << a[i].second << " "
                         << a[j].second << " "
                         << a[mid].second << "\n";
                    return 0;
                } else if (a[mid].first < need) {
                    lo = mid + 1;
                } else {
                    hi = mid - 1;
                }
            }
        }
    }

    cout << "IMPOSSIBLE\n";
    return 0;
}
CautionWatch out for integer types
  • Use long long for the values and for x, since \(x\) and \(a_i\) are up to \(10^9\), and sums can reach up to \(3 \cdot 10^9\).
  • Make sure to always search in the range [j+1, n-1] so that all three positions are distinct.