Evenly distribute arrays data between processes in MPI

I want to solve a load balancing problem using MPI with the C language. Each MPI task has an array of different size (consisting of integers).

Initial situation : data is distributed unequally between MPI processes. And we want to get arrays with same length in each process after distributing data.

We asssume each process contains an array of a radom size :

  int nbTask;
  int myRank;
  MPI_Init(&argc, &argv);

  MPI_Comm_size(MPI_COMM_WORLD, &nbTask);
  MPI_Comm_rank(MPI_COMM_WORLD, &myRank);
  time_t t;

  srand(time(NULL) + myRank);

  int size = (rand() % (60 - 40 + 1)) + 40;

I then calculate the size that each array would have after balancing :

  int global_sum;
  int new_size;

  MPI_Allreduce(&size, &global_sum, 1, MPI_INT, MPI_SUM, MPI_COMM_WORLD);
 
  new_size = (int)(((float)global_sum/nbTask) + 1);


  int exScan;
  MPI_Exscan(&size, &exScan, 1, MPI_INT, MPI_SUM, MPI_COMM_WORLD);

Can you think of an algorithm to send data from each processor to another in order to get the same size of arrays in each processor.

Appearantly, the solution uses MPI_Scan or MPI_Exscan to find what to send to processor p-1 and p+1

Here is my solution :

#include <stdlib.h>
#include <stdio.h>
#include <time.h>

#include <mpi.h>


int main(int argc, char* argv[])
{

  int nbTask;
  int myRank;

  MPI_Init(&argc, &argv);

  MPI_Comm_size(MPI_COMM_WORLD, &nbTask);
  MPI_Comm_rank(MPI_COMM_WORLD, &myRank);

  MPI_Status status;
  MPI_Request request[2];


  time_t t;

  int offset_left;
  int offset_right;

  srand(time(NULL) + myRank);

  int size = (rand() % (60 - 40 + 1)) + 40;
  printf("%d  %d  \n",myRank, size);


  int global_sum;
  int new_size;

  MPI_Allreduce(&size, &global_sum, 1, MPI_INT, MPI_SUM, MPI_COMM_WORLD);
  printf("------- global_sum : %d \n", global_sum);
  
  new_size = (global_sum + myRank) / nbTask;
  printf("task %d - new_size : %d \n", myRank, new_size);

  int new_array[new_size];


  int exScan_size;
  int exScan_newSize;

  MPI_Exscan(&size, &exScan_size, 1, MPI_INT, MPI_SUM, MPI_COMM_WORLD);
  MPI_Exscan(&new_size, &exScan_newSize, 1, MPI_INT, MPI_SUM, MPI_COMM_WORLD);

  if (myRank == 0) {
    exScan_size = 0;
    exScan_newSize = 0;
  }

  printf("task %d - exScan_size : %d \n", myRank, exScan_size);

  int array[size];
  for(int i = 0; i < size; i++) {
    array[i] = exScan_size + i;
  }

  offset_left = exScan_size - exScan_newSize;
  offset_right = offset_left + size - new_size;

  printf("task %d - offset_left : %d \n", myRank, offset_left);
  printf("task %d - offset_right : %d \n", myRank, offset_right);

  int data_left_size = abs(offset_left);
  int data_right_size = abs(offset_right);

  int data_left[data_left_size];
  int data_right[data_right_size];

  
  for(int i = 0; i < size; i++) {
    if((i + offset_left) >= 0 && (i + offset_left) < new_size) 
      new_array[i + offset_left] = array[i];
  }


  if(offset_left < 0) {
    for(int i = 0; i < data_left_size; i++) {
      data_left[i] = array[i];
    }
    if(myRank != 0)
      MPI_Isend(&(data_left[0]), data_left_size, MPI_INT, myRank-1, 0, MPI_COMM_WORLD, &request[0]);
  }
  else if (offset_left > 0) {
    MPI_Recv(&(data_left[0]), data_left_size, MPI_INT, myRank-1, 0, MPI_COMM_WORLD, &status);
     for(int i = 0; i < data_left_size; i++) {
      new_array[i] = data_left[i];
    }
  }

  if(offset_right > 0) {
    for(int i = 0; i < data_right_size; i++) {
      data_right[i] = array[size - data_right_size + i];
    }
    if(myRank != nbTask - 1)
      MPI_Isend(&(data_right[0]), data_right_size, MPI_INT, myRank+1, 0, MPI_COMM_WORLD, &request[0]);
  }
  else if (offset_right < 0) {
      MPI_Recv(&(data_right[0]), data_right_size, MPI_INT, myRank+1, 0, MPI_COMM_WORLD, &status);
      for(int i = 0; i < data_right_size; i++) {
          new_array[new_size - data_right_size + i] = data_right[i];
      }
  }

  printf("-------------------------------\n");
  printf("task %d - new_array \n", myRank);
  for(int loop = 0; loop < new_size; loop++)
      printf("%d ", new_array[loop]);
  printf("\n-------------------------------\n");

  MPI_Finalize();

  return 0;

}