USACO 2020 January Contest, Gold Problem 2. Farmer John Solves 3SUM

USACO 2020 January Contest, Gold Problem 2. Farmer John Solves 3SUM

Farmer John believes he has made a major breakthrough in algorithm design: he claims to have found a nearly linear time algorithm for the 3SUM problem, an algorithmic problem famous for the fact that no known solution exists running in substantially better than quadratic time. One formulation of the 3SUM problem is the following: given an array s1,,sms1,…,sm of integers, count the number of unordered triples of distinct indices i,j,ki,j,k such that si+sj+sk=0si+sj+sk=0.

To test Farmer John's claim, Bessie has provided an array AA of NN integers (1N50001≤N≤5000). Bessie also asks QQ queries (1Q1051≤Q≤105), each of which consists of two indices 1aibiN1≤ai≤bi≤N. For each query, Farmer John must solve the 3SUM problem on the subarray A[aibi]A[ai…bi].

Unfortunately, Farmer John has just discovered a flaw in his algorithm. He is confident that he can fix the algorithm, but in the meantime, he asks that you help him pass Bessie's test!

 

SCORING:

 

  • Test cases 2-4 satisfy N500.N≤500.
  • Test cases 5-7 satisfy N2000.N≤2000.
  • Test cases 8-15 satisfy no additional constraints.

 

 

INPUT FORMAT (file threesum.in):

The first line contains two space-separated integers NN and QQ. The second line contains the space-separated elements A1,,ANA1,…,AN of array AA. Each of the subsequent QQ lines contains two space-separated integers aiai and bibi, representing a query.It is guaranteed that 106Ai106−106≤Ai≤106 for every array element AiAi.

 

OUTPUT FORMAT (file threesum.out):

The output should consist of QQ lines, with each line ii containing a single integer---the answer to the ii-th query. Note that you should use 64-bit integers to avoid overflow.

 

SAMPLE INPUT:

7 3
2 0 -1 1 -2 3 3
1 5
2 4
1 7

SAMPLE OUTPUT:

2
1
4

For the first query, the possible triples are (A1,A2,A5)(A1,A2,A5) and (A2,A3,A4).(A2,A3,A4).

 

Problem credits: Dhruv Rohatgi

<h3> USACO 2020 January Contest, Gold Problem 2. Farmer John Solves 3SUM 题解(翰林国际教育提供,仅供参考)</h3>
<p style="text-align: center;">题解请<a href="/register" target="_blank" rel="noopener">注册</a>或<a href="/login" target="_blank" rel="noopener">登录</a>查看</p>
[/hide]

(Analysis by Benjamin Qi)

For each pair (i,k)(i,k) satisfying i<ki<k let num[i][k]num[i][k] equal the number of jj such that i<j<ki<j<k and Ai+Aj+Ak=0Ai+Aj+Ak=0. Then if ans[i][k]ans[i][k] denotes the answer for (Ai,Ai+1,,Ak)(Ai,Ai+1,…,Ak), we can write

ans[i][k]=num[i][k]+ans[i][k1]+ans[i+1][k]ans[i+1][k1].ans[i][k]=num[i][k]+ans[i][k−1]+ans[i+1][k]−ans[i+1][k−1].

After generating ansans, each query can be answered in constant time.

Now I'll describe a way to compute num[i][i+1],,num[i][N]num[i][i+1],…,num[i][N] in O(N)O(N) time. For each kk from i+1,Ni+1,…N in increasing order, consider a hashmap (unordered_map in C++) that allows you to query the number of occurrences of any integer among Ai+1,,Ak1Ai+1,…,Ak−1. Then num[i][k]num[i][k] is equal to the number of occurrences of AiAk−Ai−Ak in this map. When kk is incremented by one the only change to the map is a single insertion. As hashmap operations run in O(1)O(1) time, this solution runs in O(N2)O(N2) time overall.

However, due to the high constant factor of hashmap, this solution does not earn full points. Because all entries of AA are in the range [106,106],[−106,106], we can replace the hashmap with an array of size 2106+1,2⋅106+1, greatly improving the runtime.

 

import java.io.*;
import java.util.*;
 
public class threesum {
	public static void main(String[] args) throws IOException {
		BufferedReader in = new BufferedReader(new FileReader("threesum.in"));
		PrintWriter out = new PrintWriter(new BufferedWriter(new FileWriter("threesum.out")));
		String[] line = in.readLine().split(" ");
		int N = Integer.parseInt(line[0]);
		int Q = Integer.parseInt(line[1]);
		line = in.readLine().split(" ");
		int[] A = new int[N]; 
		long[][] ans = new long[N][N];
		for (int i = 0; i < N; ++i) A[i] = Integer.parseInt(line[i]);
		int[] z = new int[2000001];
		for (int i = N-1; i >= 0; --i) {
			for (int j = i+1; j < N; ++j) {
				int ind = 1000000-A[i]-A[j];
				if (ind >= 0 && ind <= 2000000) ans[i][j] = z[ind];
				z[1000000+A[j]] ++;
			}
			for (int j = i+1; j < N; ++j) {
				z[1000000+A[j]] --;
			}
		}
		for (int i = N-1; i >= 0; --i) 
			for (int j = i+1; j < N; ++j)
				ans[i][j] += ans[i+1][j]+ans[i][j-1]-ans[i+1][j-1];
		for (int i = 0; i < Q; ++i) {
			line = in.readLine().split(" ");
			int a = Integer.parseInt(line[0]);
			int b = Integer.parseInt(line[1]);
			out.println(ans[a-1][b-1]);
		}
		out.close();
	}
}

Of course, it was still possible to pass without replacing the hashmap by an array. Although this wasn't intended, I'll include two additional solutions for the sake of completeness.

In C++, gp_hash_table is somewhat faster than unordered_map, especially if you specify an initial capacity. See here for more information.

 

#include <bits/stdc++.h>
using namespace std;

void setIO(string name) {
    ios_base::sync_with_stdio(0); cin.tie(0);
    freopen((name+".in").c_str(),"r",stdin);
    freopen((name+".out").c_str(),"w",stdout);
}

#include <ext/pb_ds/assoc_container.hpp> // for gp_hash_table
using namespace __gnu_pbds;

int N,Q;
long long ans[5000][5000];
vector<int> A;
 
int main() {
	setIO("threesum");
	cin >> N >> Q;
	A.resize(N); for (int i = 0; i < N; ++i) cin >> A[i];
	for (int i = 0; i < N; ++i) {
		gp_hash_table<int,int> g({},{},{},{},{1<<13}); 
		// initialize with capacity that is a power of 2
		for (int j = i+1; j < N; ++j) {
			int res = -A[i]-A[j];
			auto it = g.find(res); 
			if (it != end(g)) ans[i][j] = it->second;
			g[A[j]] ++;
		}
	}
	for (int i = N-1; i >= 0; --i) for (int j = i+1; j < N; ++j)
		ans[i][j] += ans[i+1][j]+ans[i][j-1]-ans[i+1][j-1];
	for (int i = 0; i < Q; ++i) {
		int a,b; cin >> a >> b;
		cout << ans[a-1][b-1] << "\n";
	}
}

In Java, a hashmap solution passes if StreamTokenizer is used to take care of input, although it uses much more memory than I would expect. (If anyone knows how to reduce the memory usage, could you let me know?)

 

import java.io.*;
import java.util.*;
 
public class threesum {
	static StreamTokenizer in;
	static int nextInt() throws IOException {
		in.nextToken();
		return (int)in.nval;
	}
	public static void main(String[] args) throws IOException {
		in = new StreamTokenizer(new BufferedReader(new FileReader("threesum.in")));
		PrintWriter out = new PrintWriter(new BufferedWriter(new FileWriter("threesum.out")));
		int N = nextInt(); int Q = nextInt();
		int[] A = new int[N]; 
		long[][] ans = new long[N][N];
		for (int i = 0; i < N; ++i) A[i] = nextInt();
		Map<Integer,Integer> z = new HashMap<>();
		for (int i = N-1; i >= 0; --i) {
			z.clear();
			for (int j = i+1; j < N; ++j) {
				int ind = -A[i]-A[j];
				ans[i][j] = z.getOrDefault(ind,0);
				z.put(A[j],z.getOrDefault(A[j],0)+1);
			}
		}
		for (int i = N-1; i >= 0; --i) 
			for (int j = i+1; j < N; ++j)
				ans[i][j] += ans[i+1][j]+ans[i][j-1]-ans[i+1][j-1];
		for (int i = 0; i < Q; ++i) 
			out.println(ans[nextInt()-1][nextInt()-1]);
		out.close();
	}
}

[/hide]

翰林国际教育资讯二维码