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

#define N 100000000

int number_of_digits(int n)
{
    int count = 0;
    while (n > 0)
    {
        count++;
        n /= 10;
    }
    return count;
}

int num_to_array(int n, int *array)
{
    int num_of_digits = number_of_digits(n);
    int digit;

    int count = 0;
    while (n > 0)
    {
        count++;
        digit = n % 10;
        array[num_of_digits - count] = digit;
        n /= 10;
    }

    return num_of_digits;
}

int array_to_num(int *array, int length)
{
    int num = 0;
    for (int i = 0; i < length; i++)
    {
        num *= 10;
        num += array[i];
    }
    return num;
}

int kaprekarov(int n)
{
    int n2 = n * n;

    int array[10];
    int length = num_to_array(n2, array);

    int left_num, right_num;

    for (int d = length; d > 0; d--)
    {
        left_num = array_to_num(array, length - d);
        right_num = array_to_num(array + length - d, d);

        if (left_num + right_num == n)
            return 1;
    }

    return 0;
}

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

    int rank, size;

    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);

    int *kaprekarovi = (int *)malloc(N / size * sizeof(int));
    
    int start = rank + 1;
    int end = N;
    int step = size;

    int count = 0;
    for (int i = start; i < end; i += step)
        if (kaprekarov(i))
            kaprekarovi[count++] = i;

    int global_count;
    MPI_Status status;

    if (rank == 0)
    {
        global_count = count;

        for (int i = 1; i < size; i++)
        {
            MPI_Probe(i, 1234, MPI_COMM_WORLD, &status);
            MPI_Get_count(&status, MPI_INT, &count);
            MPI_Recv(kaprekarovi + global_count, count, MPI_INT, i, 1234, MPI_COMM_WORLD, &status);
            global_count += count;
        }

        for (int i = 0; i < global_count; i++)
            printf("%d\n", kaprekarovi[i]);
    }
    else
    {
        MPI_Send(kaprekarovi, count, MPI_INT, 0, 1234, MPI_COMM_WORLD);
    }

    MPI_Finalize();

    return 0;
}