histcat

Problem#

KK has a sequence of positive integers $a_1,a_2,\ldots,a_n$, and a positive integer $P$. KK considers an integer triplet $(i,j,k)$ to be good if and only if it satisfies the following conditions:

• $1 \le i < j < k \le n$;
• $P=a_i\times 2^{\lfloor\log_2 a_j\rfloor+\lfloor\log_2 a_k\rfloor+2}+a_j\times 2^{\lfloor\log_2 a_k\rfloor+1}+a_k$.

Please help KK find the number of good triplets.
For $100%$ of the data, $1\le T\le 10^3,1\le n\le 10^5,\sum n\le 10^6,1\le a_i < 2^{20},1\le P < 2^{60}$.

Solution#

First, consider a brute force approach, $O(n^3)$, which is definitely not feasible. Let's think about optimization.

We can preprocess the floor of the logarithm base 2 of each number, and then the equation can be split into factors with $k$ and without $k$, enumerating $k$ while also calculating the count of the other factor. We can store this in a map, but the time complexity $O(n^2)$ is still not feasible.

Observing the equation, we can see that it actually combines the three numbers $a_i$, $a_j$, and $a_k$ in binary to form $P$, so we enumerate $j$ and also the position of $a_j$ in $P$. During the enumeration, we maintain the count of numbers before and after $j$ in the sequence to help count the answer.

Issues#

1. $a_k$ cannot have leading zeros!!!
2. When 1 << t overflows long long, it should be changed to 1ll << t!!!

Code#

#include<bits/stdc++.h>
#define int long long
using namespace std;
int T, n;
const int N = 1e5 + 10;
int a[N];

unordered_map<int, int> suf, pre;
int len[N], L, P, bit[62];

inline int getlen(int u)
{
	int ans = 0;
	while(u)
	{
		ans ++;
		u >>= 1;
	}
	return ans;
}

signed main()
{
	bit[0] = 1;
	for(int i = 1;i <= 60;i++)
		bit[i] = bit[i - 1] * 2; 
	ios::sync_with_stdio(0), cin.tie(0);
	cin >> T;
	while(T --)
	{
		pre.clear(), suf.clear();
		int ans = 0;
		cin >> n, cin >> P, L = getlen(P);
//		cout << "P LEN" << L << endl;
		for(int i = 1;i <= n;i++)
			cin >> a[i], len[i] = getlen(a[i]);

		for(int i = 2;i <= n;i++)
			suf[a[i]]++;
		for(int j = 2;j < n;j++)
		{
			int i = j;
			pre[a[j - 1]] ++;
			suf[a[j]] --;
			for(int start = 2; start + len[i] - 1 < L;start++)
			{
				int t = (L - start - len[i] + 1);
				if(((P >> t) & ((1 << len[i]) - 1)) != a[i] || !((P >> (t - 1)) & 1)) continue;
				int before = (P >> (L - start + 1)), after = (P & ((1ll << t) - 1));
//				if(j == 3)
//					cout << before <<" " << after<< endl;
				ans += pre[before] * suf[after];
			}
//			cout << ans << endl;
		}
		cout << ans << endl;
	}
}

/*
input:
1
5 7
8 8 1 1 1

std: 1
output: 0

*/