UVa 11300 - Spreading the Wealth

Problem link: http://uva.onlinejudge.org/external/113/11300.html

Problem in short -

Given n people sitting on a circular table, each with some number of coins, you have to redistribute the coins in a manner so that everyone has equal number of coins. Any person can only pass coins to his left or right. Calculate the minimum number of coins that need to be transferred.

For example, if there were 4 people with 1, 2, 5 and 4 coins respectively, then we'll need to make sure everyone has (1+2+5+4)/4 = 3 coins finally. So, what we can do is, pass 1 coin from the last person to the first person and then third person will give 2 coins to second person who will pass 1 coin to first person. So, in total we'll need 4 transfers.

Lets define a variable x(p, p+1) = no. of coins that will be transferred from person p to person p+1. Here p+1 = (p+1) mod n. This variable can positive, negative or zero depending on which way the transfer is happening.

Now, for each person, the difference between transfers from adjacent 2 persons equals difference between his current coins and avg number of coins he should finally get [lets define it per]. Mathematically,

x(0,1) - x(1,2) = per - a[1]

Lets define, per - a[i] as A[i]. So, for n persons, we'll get -

x(0,1) = A[1] + x(1,2)
x(1,2) = A[2] + x(2,3)
.............
x(n-2, n-1) = A[n-1] + x(n-1, 0)
x(n-1, 0) = A[0] + x(0, 1)

Now, we can rewrite these n linear equations in such a way that there is always a single (and same) variable on the right hand side of the equations -

x(n-1, 0) = 0 + x(n-1, 0)
x(n-2, n-1) = A[n-1] + x(n-1, 0)
x(n-3, n-2) = A[n-2] + x(n-2, n-1) = A[n-2] + A[n-1] + x(n-1, 0)
..............
x(1,2) = A[2] + A[3] + ... + A[n-1] + x(n-1, 0)
x(0,1) = A[1] + A[2] + A[3] + ... + A[n-1] + x(n-1, 0)

As the equations tell us, we have to assign a value to x(n-1, 0) such that |x(0,1)| + |x(1,2)| + .... |x(n-1,0)| is minimized. The required value would be the median of the constants in the n equations above.

Now, to get the median, we can either sort the array first using any O(nlogn) sorting algorithm and then pick the middle element from it. Or we can use a O(n) average time, O(n^2) worst case time algorithm that is similar to randomized quick sort. You can read it further about it on Chapter 7 and 9 of Introduction to Algorithms by CLRS.

Here's the source code -

#include <stdio.h>
#include <assert.h>
#include <vector>
#include <algorithm>
#include <time.h>
#include <stdlib.h>
#include <string.h>
using namespace std;

#define MAXN 1000001
#define u64 unsigned long long
#define i64 long long
#define rep(i,n) for(i=0;i<(n);i++)
#define min(a,b) (((a)<(b))?(a):(b))
#define mabs(x) (((x)<0)?(-x):(x))

int n;
u64 a[MAXN];
i64 A[MAXN];
u64 sum, per;

int get_random(int a, int b) { //returns a random number between a & b
 int x = rand();
 int y = b - a;
 x = (int)((x / (double)RAND_MAX) * y) + a;
 if(x < a) x = a;
 if(x > b) x = b;
 return x;
}

// partition the array A[p], A[p+1] ... A[r] in such a way so that 
// all the elements left to the pivot element are smaller than it and all the elements right to it are larger than it
// return the pivot element's index
int partition(int p, int r) { 
 int x = A[r];
 int i,j;
 i = p-1;
 for(j = p; j < r; j++) {
  if(A[j] <= x) {
   i++;
   swap(A[i], A[j]);
  }
 }
 swap(A[i+1], A[r]);
 return i+1;
}

// same functionality as partition(), but selects the pivot element randomly so that no particular input can always expose the O(n^2) worst case behaviour
int randomized_partition(int p, int r) {
 int i = get_random(p, r);
 swap(A[r], A[i]);
 return partition(p, r);
}

// select the i-th element largest from the array A[p], A[p+1] ... A[r]
i64 randomized_select(int p, int r, int i) {
 if(p == r) return A[p];
 int q = randomized_partition(p, r);
 int k = q - p + 1;
 if(i == k) return A[q];
 else if(i < k) return randomized_select(p, q-1, i);
 else return randomized_select(q+1, r, i-k);
}


int main() {
 int i;
 u64 res;
 i64 mid;
 srand(time(NULL));
 while(scanf(" %d",&n) == 1) {
  sum = 0;
  rep(i,n) {
   scanf(" %llu",&a[i]);
   sum += a[i];
  }
  per = sum / n;
  rep(i,n) A[i] = (i64)per - (i64)a[i];
  for(i=n-2;i>0;i--) A[i] = A[i+1] + A[i];
  A[0] = 0;
  //sort(A, A+n);
  //mid = ( A[n/2] + A[(n-1)/2] ) / 2;
  mid = randomized_select(0, n-1, n/2 + 1);
  mid += randomized_select(0, n-1, (n-1)/2 + 1);
  mid /= 2;
  res = 0;
  rep(i,n) {
   A[i] -= mid;
   res += mabs(A[i]);
  }
  printf("%llu\n",res);
 }
 return 0;
}