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

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

    int rank;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);

    const int M = 4; // rows
    const int N = 5; // cols
    int A[M][N];

    // Initialize matrix on rank 0
    if (rank == 0) {
        for (int i = 0; i < M; i++)
            for (int j = 0; j < N; j++)
                A[i][j] = i * N + j;
    }

    // Define column type
    MPI_Datatype column_type;
    int column_index = 2; // let's send column 2
    MPI_Type_vector(M, 1, N, MPI_INT, &column_type);
    MPI_Type_commit(&column_type);

    if (rank == 0) {
        MPI_Send(&A[0][column_index], 1, column_type, 1, 0, MPI_COMM_WORLD);
    } else if (rank == 1) {
        int recv_col[M];
        MPI_Recv(recv_col, M, MPI_INT, 0, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);

        // Print received column
        for (int i = 0; i < M; i++) {
            printf("Rank 1 received col[%d] = %d\n", i, recv_col[i]);
        }
    }

    MPI_Type_free(&column_type);
    MPI_Finalize();
    return 0;
}