5 Best Ways to Sort a 2D Array Across Columns in Python

💡 Problem Formulation: The task is to take a two-dimensional array (a list of lists) in Python and sort it based on the values across columns. Consider the input [[3, 2, 1], [6, 5, 4], [9, 8, 7]]. The desired output after sorting across columns is [[1, 2, 3], [4, 5, 6], [7, 8, 9]], where each row contains sorted elements that were originally in the same column.

Method 1: Using Built-in Sort with a Custom Key Function

This method involves using the built-in sort() method in Python. A custom key function is defined to perform the sorting across columns by first transposing the array, sorting, and then transposing it back.

Here’s an example:

def sort_2d_array_columns(arr):
    transposed = list(map(list, zip(*arr)))
    for row in transposed:
        row.sort()
    return list(map(list, zip(*transposed)))

# Example array
array = [[3, 2, 1], [6, 5, 4], [9, 8, 7]]
sorted_array = sort_2d_array_columns(array)
print(sorted_array)

Output:

[[1, 2, 3], [4, 5, 6], [7, 8, 9]]

This snippet defines a function sort_2d_array_columns that takes a 2D array as an argument. It transposes the array, sorts each row (which are the original columns), and then transposes the array back to its original orientation with sorted columns.

Method 2: Using NumPy’s Sort Function

NumPy, a numerical library in Python, provides a direct way to sort arrays along any axis. Here, the sort function is used with the axis parameter to sort the array along columns.

Here’s an example:

import numpy as np

# Example array
array = np.array([[3, 2, 1], [6, 5, 4], [9, 8, 7]])
sorted_array = np.sort(array, axis=0)
print(sorted_array)

Output:

[[1 2 3]
 [4 5 6]
 [7 8 9]]

In the above code, we use NumPy’s np.sort function with the axis=0 argument, which sorts the 2D array column-wise, and the result is nicely formatted as a NumPy array.

Method 3: List Comprehension with Sorted()

This method applies the sorted() function within a list comprehension. It leverages the unpacking operator (*) to transpose the 2D array, sort each of the transposed rows, and transpose back to achieve column-wise sorting.

Here’s an example:

array = [[3, 2, 1], [6, 5, 4], [9, 8, 7]]
sorted_array = [list(i) for i in zip(*[sorted(col) for col in zip(*array)])]
print(sorted_array)

Output:

[[1, 2, 3], [4, 5, 6], [7, 8, 9]]

This snippet uses a nested list comprehension to sort the array’s columns. It transposes the array twice (before and after sorting the columns) without the need for an explicit transpose function.

Method 4: Using Pandas DataFrame

Pandas is a powerful data manipulation library. Using its DataFrame object, one can easily sort the data column-wise using the sort_values() method, specifying the axis along which to sort.

Here’s an example:

import pandas as pd

# Example array
array = [[3, 2, 1], [6, 5, 4], [9, 8, 7]]
df = pd.DataFrame(array)
sorted_array = df.apply(sorted, axis=0).values.tolist()
print(sorted_array)

Output:

[[1, 2, 3], [4, 5, 6], [7, 8, 9]]

This code snippet converts a 2D list into a Pandas DataFrame and then applies the sorted() function along each column. Finally, it converts the sorted DataFrame back into a 2D list.

Bonus One-Liner Method 5: Lambda with Sorted and Zip

A concise one-liner using lambda expressions combines the sorted function with zip to transpose and sort the array columns in one go.

Here’s an example:

array = [[3, 2, 1], [6, 5, 4], [9, 8, 7]]
sorted_array = list(map(lambda x: list(x), zip(*sorted(zip(*array), key=lambda x: x[0]))))
print(sorted_array)

Output:

[[1, 2, 3], [4, 5, 6], [7, 8, 9]]

This one-liner sorts the columns of the array inline by transposing the input list, sorting by the first element of each tuple (which are the original columns), and then transposing back to list form.

Summary/Discussion

• Method 1: Built-in Sort with Custom Key. Strengths: Utilizes Python’s built-in capabilities without the need for external libraries. Weaknesses: A bit more verbose and manually handles transposing.
• Method 2: NumPy’s Sort Function. Strengths: Very convenient and performant for numerical data. Weaknesses: Requires the NumPy library, which might be excessive for simple tasks.
• Method 3: List Comprehension with Sorted. Strengths: Concise and pythonic, uses no additional libraries. Weaknesses: Can be difficult to read and might be less efficient for large datasets.
• Method 4: Using Pandas DataFrame. Strengths: Extremely convenient for those already using pandas for data manipulation. Weaknesses: Overkill for small or simple sorting tasks.
• Bonus Method 5: Lambda with Sorted and Zip. Strengths: A one-liner solution that is neat and compact. Weaknesses: Can reduce readability and is harder to debug for beginners.