Solving 5SUM in O(n^3) with strict memory limits - c++

Question

I need a way to solve the classic 5SUM problem without hashing or with a memory efficient way of hashing.
The problem asks you to find how many subsequences in a given array of length N have the sum equal to S
Ex:
Input
6 5
1 1 1 1 1 1
Output
6
The restrictions are:
N <= 1000 ( size of the array )
S <= 400000000 ( the sum of the subsequence )
Memory usage <= 5555 kbs
Execution time 2.2s
I'm pretty sure the excepted complexity is O(N^3). Due to the memory limitations hashing doesn't provide an actual O(1) time.
The best I got was 70 points using this code. ( I got TLE on 6 tests )
#include <iostream>
#include <fstream>
#include <algorithm>
#include <vector>
#define MAX 1003
#define MOD 10472
using namespace std;
ifstream in("take5.in");
ofstream out("take5.out");
vector<pair<int, int>> has[MOD];
int v[MAX];
int pnt;
vector<pair<int, int>>::iterator it;
inline void ins(int val) {
pnt = val%MOD;
it = lower_bound(has[pnt].begin(), has[pnt].end(), make_pair(val, -1));
if(it == has[pnt].end() || it->first != val) {
has[pnt].push_back({val, 1});
sort(has[pnt].begin(), has[pnt].end());
return;
}
it->second++;
}
inline int get(int val) {
pnt = val%MOD;
it = lower_bound(has[pnt].begin(), has[pnt].end(), make_pair(val, -1));
if(it == has[pnt].end() || it->first != val)
return 0;
return it->second;
}
int main() {
int n,S;
int ach = 0;
int am = 0;
int rez = 0;
in >> n >> S;
for(int i = 1; i <= n; i++)
in >> v[i];
sort(v+1, v+n+1);
for(int i = n; i >= 1; i--) {
if(v[i] > S)
continue;
for(int j = i+1; j <= n; j++) {
if(v[i]+v[j] > S)
break;
ins(v[i]+v[j]);
}
int I = i-1;
if(S-v[I] < 0)
continue;
for(int j = 1; j <= I-1; j++) {
if(S-v[I]-v[j] < 0)
break;
for(int k = 1; k <= j-1; k++) {
if(S-v[I]-v[j]-v[k] < 0)
break;
ach = S-v[I]-v[j]-v[k];
rez += get(ach);
}
}
}
out << rez << '\n';
return 0;
}

Answer 1

I think it can be done. We are looking for all subsets of 5 items in the array arr with the correct SUM. We have array with indexes 0..N-1. Third item of those five can have index i in range 2..N-3. We cycle through all those indexes. For every index i we generate all combinations of two numbers for index in range 0..i-1 on the left of index i and all combinations of two numbers for index in the range i+1..N-1 on the right of index i. For every index i there are less than N*N combinations on the left plus on the right side. We would store only sum for every combination, so it would not be more than 1000 * 1000 * 4 = 4MB.
Now we have two sequences of numbers (the sums) and task is this: Take one number from first sequence and one number from second sequence and get sum equal to Si = SUM - arr[i]. How many combinations are there? To do it efficiently, sequences have to be sorted. Say first is sorted ascending and have numbers a, a, a, b, c ,.... Second is sorted descending and have numbers Z, Z, Y, X, W, .... If a + Z > Si then we can throw Z away, because we do not have smaller number to match. If a + Z < Si we can throw away a, because we do not have bigger number to match. And if a + Z = Si we have 2 * 3 = 6 new combinations and get rid of both a and Z. If we get sorting for free, it is nice O(N^3) algorithm.
While sorting is not for free, it is O(N * N^2 * log(N^2)) = O(N^3 * log(N)). We need to do sorting in linear time, which is not possible. Or is it? In index i+1 we can reuse sequences from index i. There are only few new combinations for i+1 - only those that involve number arr[i] together with some number from index 0..i-1. If we sort them (and we can, because there are not N*N of them, but N at most), all we need is to merge two sorted sequences. And that can be done in linear time. We can even avoid sorting completely if we sort arr at the beginning. We just merge.
For second sequence the merging does not involve adding but removing, but it is very simmilar.
The implementation seems to work, but I expect there is off by one error somewhere ;-)
#include <iostream>
#include <fstream>
#include <algorithm>
#include <vector>
using namespace std;
int Generate(int arr[], int i, int sums[], int N, int NN)
{
int p1 = 0;
for (int i1 = 0; i1 < i - 1; ++i1)
{
int ai = arr[i1];
for (int i2 = i1 + 1; i2 < i; ++i2)
{
sums[p1++] = ai + arr[i2];
}
}
sort(sums, sums + p1);
return p1;
}
int Combinations(int n, int sums[], int p1, int p2, int NN)
{
int cnt = 0;
int a = 0;
int b = NN - p2;
do
{
int state = sums[a] + sums[b] - n;
if (state > 0) { ++b; }
else if (state < 0) { ++a; }
else
{
int cnta = 0;
int lastA = sums[a];
while (a < p1 && sums[a] == lastA) { a++; cnta++; }
int cntb = 0;
int lastB = sums[b];
while (b < NN && sums[b] == lastB) { b++; cntb++; }
cnt += cnta * cntb;
}
} while (b < NN && a < p1);
return cnt;
}
int Add(int arr[], int i, int sums[], int p2, int N, int NN)
{
int ii = N - 1;
int n = arr[i];
int nn = n + arr[ii--];
int ip = NN - p2;
int newP2 = p2 + N - i - 1;
for (int p = NN - newP2; p < NN; ++p)
{
if (ip < NN && (ii < i || sums[ip] > nn))
{
sums[p] = sums[ip++];
}
else
{
sums[p] = nn;
nn = n + arr[ii--];
}
}
return newP2;
}
int Remove(int arr[], int i, int sums[], int p1)
{
int ii = 0;
int n = arr[i];
int nn = n + arr[ii++];
int pp = 0;
int p = 0;
for (; p < p1 - i; ++p)
{
while (ii <= i && sums[pp] == nn)
{
++pp;
nn = n + arr[ii++];
}
sums[p] = sums[pp++];
}
return p;
}
int main() {
ifstream in("take5.in");
ofstream out("take5.out");
int N, SUM;
in >> N >> SUM;
int* arr = new int[N];
for (int i = 0; i < N; i++)
in >> arr[i];
sort(arr, arr + N);
int NN = (N - 3) * (N - 4) / 2 + 1;
int* sums = new int[NN];
int combinations = 0;
int p1 = 0;
int p2 = 1;
for (int i = N - 3; i >= 2; --i)
{
if (p1 == 0)
{
p1 = Generate(arr, i, sums, N, NN);
sums[NN - 1] = arr[N - 1] + arr[N - 2];
}
else
{
p1 = Remove(arr, i, sums, p1);
p2 = Add(arr, i + 1, sums, p2, N, NN);
}
combinations += Combinations(SUM - arr[i], sums, p1, p2, NN);
}
out << combinations << '\n';
return 0;
}