A plant#12196

A plant (Editorial)

We have to count number of pair of element such that $(a_{i} + a_{j})^{2} = a_{i} + a_{j}$ (a_i + a_j) ^ 2 = a_i + a_j

$(a_{i} + a_{j})^{2} = a_{i} + a_{j}$ (a_i + a_j) ^ 2 = a_i + a_j $\Rightarrow$ ${a_{i} + a_{j} = 1 a_{i} + a_{j} = 0$ \left\{ \begin{array}{rcl} a_i + a_j = 1 \\ a_i + a_j = 0 \end{array}\right. $\Rightarrow$ ${a_{i} = 1 - a_{j} a_{i} = - a_{j}$ \left\{ \begin{array}{rcl} a_i = 1 - a_j \\ a_i = -a_j \end{array}\right.

Solution fill be, we bruteforce index j and then we have to canculate number of such index i that satisfy $a_{i} = 1 - a_{j}$ a_i = 1 - a_j or $a_{i} = - a_{j}$ a_i = -a_j.We can do it using map

signed main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);
    cout.tie(nullptr);
    int n;
    cin >> n;
    map<int, int> cnt;

    long long ans = 0;
    for (int i = 0; i < n; i++) {
        int x;
        cin >> x;

        ans += cnt[-x];
        ans += cnt[-x+1];
        cnt[x]++;
    }

    cout << ans;
}