← All lessons

Sort, argsort, and top-n

sort, argsort, argmax, argmin, and partition.

sort orders values. argsort orders positions so you can rank rows of a table you did not flatten.

Goal

Sort a vector, rank cities with argsort / argmax, and take a cheap top-n with partition.

Sort a copy vs in place

rain = np.array([50, 40, 80, 150], dtype=float)
print("sorted copy:", np.sort(rain))
print("original:", rain)
rain.sort()
print("after rain.sort():", rain)

np.sort returns a new array. arr.sort() mutates.

Rank cities

# Nairobi, Mombasa, Kisumu, Nakuru, Eldoret — Apr mm
apr = np.array([150, 90, 180, 120, 110], dtype=float)
order = np.argsort(apr)
print("driest → wettest indices:", order)
print("driest → wettest values:", apr[order])
print("wettest index:", np.argmax(apr), "value:", apr.max())

Kisumu (index 2) is wettest.

Sort each row

rain = np.array(
    [
        [50, 40, 80, 150],
        [20, 15, 30, 90],
        [70, 80, 120, 180],
    ],
    dtype=float,
)
print(np.sort(rain, axis=1))

Months per city, low to high. The original month order is lost — that is what argsort is for when you still need labels.

partition for top-n

apr = np.array([150, 90, 180, 120, 110], dtype=float)
part = np.partition(apr, -2)
print(part)
print("two wettest (unordered):", part[-2:])

partition is cheaper than a full sort when you only need “the largest two”, not their order.

Pitfall

a.sort(axis=0) reorders down each column independently. Rows stop being “one city”. Prefer np.argsort on a 1-D summary instead.