KOL1506 - Editorial



Author: Balajiganapathi
Tester: Kevin Atienza
Editorialist: Kevin Atienza


Binomial theorem, dynamic programming


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?


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:

C_{i+1,a} = C_{i,a} + (-A_i)^{k-a}

The answer is now:

2\sum_{i=1}^N \sum_{a=0}^k {k \choose a} C_{i,a} A_i^a

By precomputing binomial coefficients and powers of A_i, this answer can be computed in O(Nk) time.


Based on the problem itself, the answer is the following sum:

\sum_{\substack{1 \le i, j \le N\\\ j \not= i}} |A_i - A_j|^k

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.

2\sum_{1 \le j < i \le N} |A_i - A_j|^k

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:

2\sum_{1 \le j < i \le N} (A_i - A_j)^k


2\sum_{i=1}^N \sum_{j=1}^{i-1} (A_i - A_j)^k

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:

\begin{aligned} \sum_{j=1}^{i-1} (A_i - A_j)^k &= \sum_{j=1}^{i-1} \sum_{a=0}^k {k \choose a} A_i^a (-A_j)^{k-a} \\\ &= \sum_{a=0}^k {k \choose a} \left(\sum_{j=1}^{i-1} (-A_j)^{k-a}\right) A_i^a\\\ \end{aligned}

Let’s define C_{i,a} = \sum_{j=1}^{i-1} (-A_j)^{k-a} so that the sum above is simply equal to

\sum_{a=0}^k {k \choose a} C_{i,a} A_i^a

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:

C_{i+1,a} = C_{i,a} + (-A_i)^{k-a}

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)



@admin The problems are still not moved to practice section!

@kevinsogo Thank you for taking time to write such nice and detailed editorials, these help a lot! :slight_smile:

1 Like

whats wrong with my code


int main()


int t;




int n,k,i,j;

long long int A[100001],cost=0,d;

scanf("%d %d",&n,&k);













//int price=cost%1000000007;



return 0;


1 Like

Hello @swamicoder,

were you able to figure out what is wrong with your code. I am implementing similar algorithm. This may be a naive approach for solving such kind of problem, but I am unable to understand why I am getting Wrong Answer. If you have figured out why this is happening, please let me know.


@admin setter’s and tester’s solutions are not available, please fix it…

1 Like

Below is my logic. (Brute Force)(Java)

private static Long processInput(TestCase t) {
		Long result = (long) 0;
		Long currResult = (long) 0;
		List<Integer> values = t.getValue();
		for(int i=0; i< t.getSize(); i++){
			for(int j=i+1; j<t.getSize(); j++) {
				currResult = (long)Math.pow(Math.abs(values.get(i) - values.get(j)), t.getPower());
				currResult = currResult % 1000000007;
				result = result + currResult;
				result = result % 1000000007;
		result *= 2;
		return result % 1000000007;

Sample Input mentioned in the problem page is giving expected output.
But on submitting, I get a WA.

Any help??

Can some one please provide me with more test cases ?
The program I wrote for this problem is working fine on my system for the given test cases. But when I submit it, I get the response as “wrong answer”. Any help with test cases or explanation will be much appreciated. Thanks in advance.

I have the same issue.
My program works for the given test cases properly…but gives wrong answer on submission
And also another doubt…they have mentioned that two houses can be at the same position.
So what happens then? Do we have to calculate the amount for that too?

The above solutions are giving WA because of the inbuilt pow function used.
The ranges of ai=[0,1e9+6] and k=[0,100], so maximum power which needs to be computed is pow(1e9+6-0,100-0), clearly this value cannot be handled by this function. So, I built a recursive power function which calculates power and takes mod at the same time.
ll power(ll a,ll b){
if(b==0) return 1;
if(b==1) return a;
ll p=power(a,ll(b/2));
if(b%2) return (p
return p;

But that solution is giving me TLE.