Maximize the sum of Kth column of a Matrix

Given two integers N and K, the task is to maximize the sum of the K^th column of N * N row-wise sorted matrix consisting of element in the range [1, N²].

Examples:

Input: N = 2, K = 2
Output: {{1, 3}, {2, 4}}
Explanations: The possible row-wise sorted matrices are [{{1, 2}, {3, 4}}, {{1, 3}, {2, 4}}, {{1, 4}, {2, 3}}, {{3, 4}, {1, 2}}, {{2, 4}, {1, 3}}, {{2, 3}, {1, 4}} ] 
Out of all the above possible matrices, the matrices [{{1, 3}, {2, 4}}, {{2, 4}, {1, 3}}, {{1, 4}, {2, 3}}, {{2, 3}, {1, 4}}] contains the maximum possible sum of the K^th column. 
Therefore, one of the possible output is {{1, 3}, {2, 4}}.

Input: N = 3, K = 2
Output: {{1, 4, 5}, {2, 6, 7}, {3, 8, 9}}

Approach: The idea here is to first fill the indices smaller than the K^th columns of the matrix by the values from the range [1, N * (K - 1)] and then fill all the elements at columns greater than or equal to the k^th column by values from the range [N * (K - 1) + 1, N²] as shown in the image below.

Follow the steps below to solve the problem:

1. Fill all the columns of the matrix smaller than K by the values from the range [1, N * (K - 1)].
2. Then, fill the columns of the matrix greater than or equal to K by the values from the range [N * (K - 1) + 1, N * N].
3. Finally, print the matrix.

Below is the implementation of the above approach:

// C++ program to implement
// the above approach

#include <bits/stdc++.h>
using namespace std;

// Function to maximize the Kth column sum
int** findMatrix(int N, int K)
{

    // Store all the elements of the
    // resultant matrix of size N*N
    int** mat = (int**)malloc(
        N * sizeof(int*));

    for (int i = 0; i < N; ++i) {
        mat[i] = (int*)malloc(
            N * sizeof(int));
    }

    // Store value of each
    // elements of the matrix
    int element = 1;

    // Fill all the columns < K
    for (int i = 0; i < N; ++i) {

        for (int j = 0; j < K - 1; ++j) {
            mat[i][j] = element++;
        }
    }

    // Fill all the columns >= K
    for (int i = 0; i < N; ++i) {

        for (int j = K - 1; j < N; ++j) {
            mat[i][j] = element++;
        }
    }

    return mat;
}

// Function to print the matrix
void printMatrix(int** mat, int N)
{

    for (int i = 0; i < N; ++i) {
        for (int j = 0; j < N; ++j) {
            cout << mat[i][j] << " ";
        }
        cout << endl;
    }
}

// Driver Code
int main()
{

    int N = 3, K = 2;
    int** mat = findMatrix(N, K);

    printMatrix(mat, N);
}

// Java program to implement
// the above approach
class GFG{

// Function to maximize the Kth column sum
static int [][]findMatrix(int N, int K)
{

    // Store all the elements of the
    // resultant matrix of size N*N
    int [][]mat = new int[N][N];

    // Store value of each
    // elements of the matrix
    int element = 1;

    // Fill all the columns < K
    for(int i = 0; i < N; ++i)
    {
        for(int j = 0; j < K - 1; ++j)
        {
            mat[i][j] = element++;
        }
    }

    // Fill all the columns >= K
    for(int i = 0; i < N; ++i) 
    {
        for(int j = K - 1; j < N; ++j)
        {
            mat[i][j] = element++;
        }
    }
    return mat;
}

// Function to print the matrix
static void printMatrix(int [][]mat, int N)
{
    for(int i = 0; i < N; ++i) 
    {
        for(int j = 0; j < N; ++j)
        {
            System.out.print(mat[i][j] + " ");
        }
        System.out.println();
    }
}

// Driver Code
public static void main(String[] args)
{
    int N = 3, K = 2;
    int [][]mat = findMatrix(N, K);

    printMatrix(mat, N);
}
}

// This code is contributed by Amit Katiyar 

# Python3 program to implement
# the above approach

# Function to maximize the Kth 
# column sum
def findMatrix(N, K):
    
    # Store all the elements of the
    # resultant matrix of size N*N
    mat = [[0 for i in range(N)]
              for j in range(N)];

    # Store value of each
    # elements of the matrix
    element = 0;

    # Fill all the columns < K
    for i in range(0, N):
        for j in range(0, K - 1):
            element += 1;
            mat[i][j] = element;

    # Fill all the columns >= K
    for i in range(0, N):
        for j in range(K - 1, N):
            element += 1;
            mat[i][j] = element;

    return mat;

# Function to print the matrix
def printMatrix(mat, N):
    
    for i in range(0, N):
        for j in range(0, N):
            print(mat[i][j], end = " ");

        print();

# Driver Code
if __name__ == '__main__':
    
    N = 3; K = 2;
    mat = findMatrix(N, K);

    printMatrix(mat, N);

# This code is contributed by Amit Katiyar 

// C# program to implement
// the above approach
using System;

class GFG{

// Function to maximize the Kth column sum
static int [,]findMatrix(int N, int K)
{

    // Store all the elements of the
    // resultant matrix of size N*N
    int [,]mat = new int[N, N];

    // Store value of each
    // elements of the matrix
    int element = 1;

    // Fill all the columns < K
    for(int i = 0; i < N; ++i)
    {
        for(int j = 0; j < K - 1; ++j)
        {
            mat[i, j] = element++;
        }
    }

    // Fill all the columns >= K
    for(int i = 0; i < N; ++i) 
    {
        for(int j = K - 1; j < N; ++j)
        {
            mat[i, j] = element++;
        }
    }
    return mat;
}

// Function to print the matrix
static void printMatrix(int [,]mat, int N)
{
    for(int i = 0; i < N; ++i) 
    {
        for(int j = 0; j < N; ++j)
        {
            Console.Write(mat[i, j] + " ");
        }
        Console.WriteLine();
    }
}

// Driver Code
public static void Main(String[] args)
{
    int N = 3, K = 2;
    int [,]mat = findMatrix(N, K);

    printMatrix(mat, N);
}
}

// This code is contributed by Amit Katiyar 

<script>

// JavaScript program to implement
// the above approach

// Function to maximize the Kth column sum
function findMatrix(N, K)
{
    
    // Store all the elements of the
    // resultant matrix of size N*N
    let mat =  new Array(N);

    // Loop to create 2D array using 1D array
    for(var i = 0; i < mat.length; i++)
    {
        mat[i] = new Array(2);
    }
  
    // Store value of each
    // elements of the matrix
    let element = 1;
  
    // Fill all the columns < K
    for(let i = 0; i < N; ++i)
    {
        for(let j = 0; j < K - 1; ++j)
        {
            mat[i][j] = element++;
        }
    }
  
    // Fill all the columns >= K
    for(let i = 0; i < N; ++i) 
    {
        for(let j = K - 1; j < N; ++j)
        {
            mat[i][j] = element++;
        }
    }
    return mat;
}
  
// Function to print the matrix
function printMatrix(mat, N)
{
    for(let i = 0; i < N; ++i) 
    {
        for(let j = 0; j < N; ++j)
        {
            document.write(mat[i][j] + " ");
        }
        document.write("<br/>");
    }
}

// Driver code
let N = 3, K = 2;
let mat = findMatrix(N, K);

printMatrix(mat, N);

// This code is contributed by souravghosh0416

</script>

1 4 5 
2 6 7 
3 8 9

Time Complexity: O(N²)
Auxiliary Space: O(N²)