位运算—— AND, OR and square sum

$deception$
Gottfried learned about binary number representation. He then came up with this task and presented it to you.

You are given a collection of n non-negative integers $a_1,\dots ,a_n$. You are allowed to perform the following operation: choose two distinct indices $1\le i,j\le n$. If before the operation $a_i=x,a_j=y$, then after the operation $a_i=xANDy$, $a_j=xORy,$ where AND and OR are bitwise AND and OR respectively (refer to the Notes section for formal description). The operation may be performed any number of times (possibly zero).

After all operations are done, compute ${\sum }_{i=1}^n{a}_i^2$ — the sum of squares of all $a_i$. What is the largest sum of squares you can achieve?

$Input$
The first line contains a single integer $n$ $(1\le n\le 2\cdot 10^5)$.

The second line contains n integers $a_1,\dots ,a_n(0\le a_i<2^{20})$.

$Output$
Print a single integer — the largest possible sum of squares that can be achieved after several (possibly zero) operations.

比如我有两个数$x$,$y$，他们的二进制位中1的个数总共有$n$个，然后我对x和y分别做与和或的操作，他们的二进制位中1的总数是不会变的。所以这题我们只要记录每个位上1的个数，然后贪心让数最大分配即可。

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
const int N = 2e5 + 10;
const int INF = 0x3f3f3f3f;
const double eps = 1e-7;
#define endl '\n'
#define pb(a) push_back(a)
#define ALL(x) x.begin(), x.end()
#define SIZE(x) int(x.size())
#define IOS ios::sync_with_stdio(0);
int n, a[N];
int bit[25];
int main() {
    scanf("%d", &n);
    for (int i = 1; i <= n; ++i) {
        scanf("%d", a + i);
        //记录1的个数
        for (int j = 0; j < 20; ++j) {
            bit[j] += (a[i] & 1);
            a[i] >>= 1;
        }
    }

    ll ans = 0;
    for (int i = 1; i <= n; ++i) {
        int x = 0;
        //开始分配，尽可能让数大
        for (int j = 0; j < 20; ++j) {
            if (bit[j]) {
                x += (1 << j);
                bit[j]--;
            }
        }
        ans += 1ll * x * x;
    }
    cout << ans << endl;

    return 0;
}