PROBLEM LINK:
Author: Balajiganapathi
Tester: Kevin Atienza
Editorialist: Kevin Atienza
PREREQUISITES:
Binomial theorem, dynamic programming
PROBLEM:
There are N houses in a row. The position of the i th house is A_i. Each house sends a gift to every other house, and the cost of sending a gift from a house at position x to a house at position y is |x-y|^k. What is the total amount of money paid by all houses, modulo 10^9 + 7?
QUICK EXPLANATION:
For 0 \le i \le N and 0 \le a \le k, let C_{i,a} := \sum_{j=1}^{i-1} (-A_j)^{k-a}. These values can be computed in O(Nk) using the recurrence:
The answer is now:
By precomputing binomial coefficients and powers of A_i, this answer can be computed in O(Nk) time.
EXPLANATION:
Based on the problem itself, the answer is the following sum:
We can reduce this sum in half by noticing that the cost of sending a gift from house i to house j is the same as sending a gift from house j to house i.
To remove the nasty absolute value in the formula, we can just sort the $A_i$s, so that if j < i, then A_j \le A_i. So assuming the $A_i$s are sorted, the sum is now:
or
For a fixed i, 1 \le i \le N, we can interpret the inner sum as the total cost of sending a gift from house i to all the other houses j < i. If we want to compute the answer quickly, we must be able to compute the inner sum quickly for each i.
We can manipulate the inner sum with the binomial theorem:
Let’s define C_{i,a} = \sum_{j=1}^{i-1} (-A_j)^{k-a} so that the sum above is simply equal to
Computing this sum for all i naĂŻvely will take O(N^2k) time, because computing C_{i,a} takes linear time for each i. We must be able to compute the $C_{i,a}$s quicker than this.
We can take advantage of the fact that C_{i,a} and C_{i+1,a} are related; The value C_{i+1,a} - C_{i,a} is just (-A_i)^{k-a}. Thus, we have the following simple recurrence:
Thus, if we know C_{i,a} for all a, then we can use this recurrence to compute C_{i+1,a} for all a quickly, and if we can compute the C_{i,a} values quickly, then we can find the answer to the problem quickly!
The pseudocode is as follows:
mod = 10^9 + 7
C = [0,0,0,...0] # length k+1
P = [0,0,0,...0] # length k+1
for i = 1...N:
// compute powers of A[i]
power = 1
for a = 0...k:
P[a] = power
power *= A[i]
// adjust answer
for a = 0...k:
answer += binom(k,a) * C[a] * P[a]
// compute powers of -A[i]
power = 1
for a = 0...k:
P[a] = power
power *= -A[i]
// adjust C's
for a = 0...k:
C[a] += P[k-a]
The running time is clearly O(Nk), assuming we have already precomputed the binomial coefficients. Don’t forget to reduce intermediate values modulo 10^9 + 7!
Time Complexity:
O(N \log N + Nk)