A gift with a purchase
A shop is running a promotion: a customer chooses one item and one gift, and the set must cost exactly s tenge, not a tenge more or less. The marketers try different sums s and for each want to know in how many ways a set can be put together. There are many items and gifts, and many questions too.
Two ways are different if the item or the gift differs, even when their prices are the same.
Input
The first line holds three numbers n, m, q (1 ≤ n, m, q ≤ 10⁵): the number of items, gifts and questions. The second line holds the n prices of the items, the third the m prices of the gifts; all prices are integers from 1 to 10⁵. Then come q lines with one number s each (0 ≤ s ≤ 2 · 10⁵).
Output
For each question, the number of “item, gift” pairs whose prices add up to exactly s.
Examples
Example 1
4 3 5 1500 2000 2000 3500 500 1000 1500 3000 3500 2500 100 5000
3 2 3 0 1
Limits
- Time · Python
- 5000 ms
- Time · Go
- 2000 ms
- Memory
- 256 MB
- Judging
- 6 tests
Walkthrough
Open it step by step: each step gives away a little more. You can jump straight to the solution, but a problem solved with one hint stays with you longer.
1Hint 1
Going over the pairs is 10¹⁰ operations. Going over the distinct item prices for each question is as many in the worst case. Can all possible s be answered at once?
2Hint 2
Let A(x) = Σ (number of items at price p) · xp and B(x) the same for gifts. What is the coefficient of xs in the product A(x) · B(x)?
3The idea
The coefficient of xs in A · B is the sum of the products ap · bs−p over all p, that is, exactly the number of pairs with sum s. So the answers to all the questions are the coefficients of the product of two polynomials of degree up to 10⁵, and the fast Fourier transform computes it in O(C log C), where C is the largest price.
4Walkthrough
Build count arrays a and b of length 10⁵ + 1, pad them with zeros to a power of two of at least 2 · 10⁵ + 1, otherwise the circular convolution will wrap large sums onto small ones, do two forward transforms, multiply pointwise, do the inverse transform and round to the nearest integer. In Python this is five lines with numpy.fft.rfft and irfft; numpy.convolve will not do, since it convolves head-on, in quadratic time.
The Go solution does one transform instead of two: the items go into the real part, the gifts into the imaginary part. The square of a + ib is a² − b² + 2iab, and the imaginary part of the square of the polynomial, divided by two, is exactly the convolution of a and b.
Precision: the answer goes up to 10¹⁰ (the third test is a hundred thousand identical items and a hundred thousand identical gifts), while the rounding error in an FFT of this length is many orders of magnitude below 0.5, so rounding recovers the integer exactly.
4+Solution
5Complexity
The convolution is O(C log C) with C = 2¹⁸, an answer to a question O(1). O(n + m + q + C log C) in all.
6Common mistakes
- A loop over prices for each question. Correct, but O(q · D) with D distinct prices, up to 10¹⁰ steps; the fourth test does not let it through.
- A circular convolution without zero padding. The sum 150,000 with length 2¹⁷ = 131,072 “wraps round” and is added to the sum 18,928.
- A 32-bit answer. 10¹⁰ does not fit in an int32.
- Dropping the fractional part instead of rounding. 2.9999999 after an inverse FFT is 3, not 2.