MPI communicator creation exercise

I'm working on an assignment which entails creating communicators following the diagonals of an n x n process grid. I'm interested to get feedback regarding the correctness of the solution I developed as well as feedback regarding the efficiency of the solution.

```
#include
#include
#include
#include

/* Example Matrix
* 0 1 2 3
* 4 5 6 7
* 8 9 10 11
* 12 13 14 15
*/

// Determine what the heads are.
// "Heads" are the first (top) value in the diagonal.
int find_heads(int world_size, int heads, int root){
// Get the "regular" heads. (0, 1, 2, 3, 4 in the Example Matrix)
int i;
int head = -1;
printf("The values in the 'heads' array represents the top values in the diagonal.\n");
for(i = 0; i <= root; i++){
head++;
heads[i] = head;
printf("heads[%d] = %d\n", i, heads[i]);
}

// Get the "irregular" (8, 12 in the Example Matrix) heads.
int j;
for(j = 0; j < (root - 2); j++){
head += root;
heads[i] = head;
printf("heads[%d] = %d\n", i, heads[i]);
i++;
}

return heads;
}

// See if rank is in the heads array.
int head_check(int rank, int *heads, int heads_length){
int i;
for( i = 0; i < heads_length; i++){
if (heads[i] == rank){
return i;
}
}
return -1;
}

// Determine what "color" a rank should be in.
// A color defines a which communicator a rank will be in.
// In this case we are grouping on the diagonals.
int find_color(int world_size, int rank, int *heads, int heads_length){
int color;
// How much to subract from a rank until it eventually matches one of the heads.
int sub = sqrt(world_size) + 1;
color = -1;
// Keep subracting from the rank until it matches one of the heads.
// If it matches a head it is in the same diagonal as that head.
while(color == -1){
color = head_check(rank, heads, heads_length);
rank -= sub;
}

return color;
}

int main(int arg

Can be simplified

Currently, you create a heads array and then use two functions find_color() and head_check() to find the color of each element using the array. But actually, you can determine the color of each element in a much simpler fashion, without ever needing to create a temporary array in the first place.

Essentially, you are trying to trace back each element to its diagonal starting point ("head"), where the starting point may either be on the top or left edge of the matrix. You can determine which edge a given element's head is on by comparing its row number against its column number. Any element where col >= row will have its head be on the top edge, otherwise its head will be on the left edge.

Given the above, you can reduce all of your color finding code to this one function:

// Rank is the element number, in the range 0..(size*size-1)
// Size is the length of the matrix on one side.
int find_color(int rank, int size)
{
    int row = rank / size;
    int col = rank % size;

    if (col >= row) {
        // Diagonal head must be in top row.
        return col - row;
    } else {
        // Diagonal head must be on left side.
        return (row - col) + size - 1;
    }
}

StackExchange Code Review Q#141821, answer score: 5