#include <stdio.h>
#include <mpi.h>

int main(void) {
    int my_rank, comm_sz;
    MPI_Init(NULL, NULL);
    MPI_Comm_rank(MPI_COMM_WORLD, &my_rank);
    MPI_Comm_size(MPI_COMM_WORLD, &comm_sz);

    const int n = 12;
    if (n % comm_sz != 0) {
        if (my_rank == 0) {
            printf("Use 1, 2, 3, 4, 6, or 12 processes.\n");
        }
        MPI_Finalize();
        return 1;
    }

    int local_n = n / comm_sz;
    int numbers[12];
    int squares[12];
    int local_numbers[12];
    int local_squares[12];

    if (my_rank == 0) {
        for (int i = 0; i < n; i++) {
            numbers[i] = i + 1;
        }
    }

    MPI_Scatter(numbers, local_n, MPI_INT,
                local_numbers, local_n, MPI_INT,
                0, MPI_COMM_WORLD);

    for (int i = 0; i < local_n; i++) {
        local_squares[i] = local_numbers[i] * local_numbers[i];
    }

    MPI_Gather(local_squares, local_n, MPI_INT,
               squares, local_n, MPI_INT,
               0, MPI_COMM_WORLD);

    if (my_rank == 0) {
        printf("Squares:");
        for (int i = 0; i < n; i++) {
            printf(" %d", squares[i]);
        }
        printf("\n");
    }

    MPI_Finalize();
    return 0;
}