Sorting, Searching, Unique Values and Set Operations

This lesson answers everyday questions about data: "put these in order", "who are the top 3?", "which grade band does each mark fall into?", "which cities appear and how often?", and "which students are in both classes?".

How to run the examples

Create a new file named sorting_searching_sets.py in your NumpyTutorials folder (see the setup in lesson 01). Any name works; this one just matches the topic. Start the file with these lines:

import numpy as np

Then copy the examples of this page into the same file, one after another, in the order shown, and run it with python sorting_searching_sets.py (with your virtual environment active). Some examples use variables created in earlier ones, so keep everything in the one file. Each Output box shows what that example prints.

Most examples use these two arrays, so add them right after the import line:

names = np.array(["Asha", "Ravi", "Neha", "Kiran", "Meena", "Arjun"])
marks = np.array([82, 34, 67, 91, 28, 67])

1. Sorting

print("np.sort(marks)   :", np.sort(marks))
print("marks unchanged  :", marks)
print("descending       :", np.sort(marks)[::-1])
print("names A-Z        :", np.sort(names))

copy = marks.copy()
copy.sort()                                  # in place, returns None
print("copy.sort()      :", copy)

np.sort(marks)   : [28 34 67 67 82 91]
marks unchanged  : [82 34 67 91 28 67]
descending       : [91 82 67 67 34 28]
names A-Z        : ['Arjun' 'Asha' 'Kiran' 'Meena' 'Neha' 'Ravi']
copy.sort()      : [28 34 67 67 82 91]

• np.sort(arr) returns a sorted copy; the original stays the same.
• arr.sort() sorts the array itself and returns None. Don't write x = arr.sort().
• There is no "descending" option: sort, then reverse with [::-1].
• Text sorts alphabetically.

2. Sorting a 2D array

table = np.array([[9, 2, 7],
                  [4, 8, 1]])
print("sort each row (axis=1, default):\n", np.sort(table))
print("sort each column (axis=0):\n", np.sort(table, axis=0))
print("sort everything (axis=None):", np.sort(table, axis=None))

sort each row (axis=1, default):
 [[2 7 9]
 [1 4 8]]
sort each column (axis=0):
 [[4 2 1]
 [9 8 7]]
sort everything (axis=None): [1 2 4 7 8 9]

By default each row is sorted on its own (axis=-1, the last axis). axis=0 sorts each column, and axis=None flattens and sorts everything. Note that sorting rows or columns separately breaks the link between values in the same row; to keep rows together, use argsort.

3. argsort: sort one array by another

argsort does not return sorted values; it returns the order of the indexes that would sort the array. That order can then be used on any related array with fancy indexing.

order = np.argsort(marks)
print("argsort(marks)       :", order)
print("names, lowest first  :", names[order])
print("names, highest first :", names[order[::-1]])
print("top 3                :", names[np.argsort(-marks)[:3]])

# Sort by marks (high to low), and by name when marks are equal:
by_marks_then_name = np.lexsort((names, -marks))     # LAST key is the main one
print("by marks, then name  :", list(zip(names[by_marks_then_name].tolist(), marks[by_marks_then_name].tolist())))

argsort(marks)       : [4 1 5 2 0 3]
names, lowest first  : ['Meena' 'Ravi' 'Arjun' 'Neha' 'Asha' 'Kiran']
names, highest first : ['Kiran' 'Asha' 'Neha' 'Arjun' 'Ravi' 'Meena']
top 3                : ['Kiran' 'Asha' 'Neha']
by marks, then name  : [('Kiran', 91), ('Asha', 82), ('Arjun', 67), ('Neha', 67), ('Ravi', 34), ('Meena', 28)]

• argsort(marks) = [4 1 5 2 0 3]: index 4 (28) is the smallest, then index 1 (34), and so on.
• names[order] lists the names in the same order as the marks.
• np.argsort(-marks) sorts by the negative, so it gives a high-to-low order; [:3] keeps the top 3.
• np.lexsort sorts by several keys. The last key is the main one: by marks high to low, and by name when the marks are equal (Arjun before Neha, both 67).

4. Searching

print("where marks > 60  :", np.where(marks > 60)[0])
print("nonzero(marks > 60):", np.nonzero(marks > 60)[0], "(same thing)")
print("first index of 67 :", np.argmax(marks == 67))
print("is 91 in marks?   :", 91 in marks)

where marks > 60  : [0 2 3 5]
nonzero(marks > 60): [0 2 3 5] (same thing)
first index of 67 : 2
is 91 in marks?   : True

• np.where(condition)[0] and np.nonzero(condition)[0] give the positions where the condition is True.
• np.argmax(marks == 67) finds the first position of a value: the first True in the mask.
• The Python in operator works on arrays too.

5. searchsorted: grade bands

np.searchsorted(sorted_array, values) tells you where each value would be inserted to keep the array sorted. That makes it a fast way to put values into bands.

limits = np.array([35, 50, 75, 90])             # must be sorted
labels = np.array(["Fail", "C", "B", "A", "A+"])
positions = np.searchsorted(limits, marks, side="right")
print("marks    :", marks)
print("positions:", positions)
print("grades   :", labels[positions])
print("insert 70 into sorted [10, 40, 80] at index:", np.searchsorted([10, 40, 80], 70))

marks    : [82 34 67 91 28 67]
positions: [3 0 2 4 0 2]
grades   : ['A' 'Fail' 'B' 'A+' 'Fail' 'B']
insert 70 into sorted [10, 40, 80] at index: 2

The limits are 35, 50, 75 and 90. A mark of 82 would go after 35, 50 and 75, which is position 3, so labels[3] = "A". A mark of 28 goes before everything (position 0), so it gets "Fail". side="right" makes a mark that equals a limit (such as exactly 50) go into the higher band. The array you search in must be sorted.

6. partition and argpartition

To find the top 3 out of a million values, you don't need to sort all million. np.partition(arr, -3) only makes sure the 3 largest values end up in the last 3 places, which is faster than a full sort.

big = np.array([15, 3, 99, 42, 8, 77, 61, 23])
# partition only promises that the 3 largest values end up in the last 3 places;
# the order of the other values is not guaranteed.
print("partition(big, -3)[-3:]:", np.partition(big, -3)[-3:], "(the 3 largest)")
top3 = np.argpartition(big, -3)[-3:]
print("3 largest values  :", np.sort(big[top3])[::-1])

partition(big, -3)[-3:]: [61 77 99] (the 3 largest)
3 largest values  : [99 77 61]

argpartition gives the positions instead of the values, just like argsort. The 3 values come out in no particular order, so we sort just those 3 at the end.

7. np.unique

cities = np.array(["Pune", "Mumbai", "Pune", "Delhi", "Mumbai", "Pune"])

print("unique              :", np.unique(cities))
values, counts = np.unique(cities, return_counts=True)
print("counts              :", dict(zip(values.tolist(), counts.tolist())))
print("most common         :", values[counts.argmax()])

values, first_index = np.unique(cities, return_index=True)
print("first position      :", dict(zip(values.tolist(), first_index.tolist())))

values, codes = np.unique(cities, return_inverse=True)
print("city codes          :", codes, "->", values)

unique              : ['Delhi' 'Mumbai' 'Pune']
counts              : {'Delhi': 1, 'Mumbai': 2, 'Pune': 3}
most common         : Pune
first position      : {'Delhi': 3, 'Mumbai': 1, 'Pune': 0}
city codes          : [2 1 2 0 1 2] -> ['Delhi' 'Mumbai' 'Pune']

8. Set operations

python_class = np.array(["Asha", "Ravi", "Neha", "Kiran"])
sql_class = np.array(["Neha", "Kiran", "Meena", "Arjun"])

print("both classes (intersect1d):", np.intersect1d(python_class, sql_class))
print("any class (union1d)       :", np.union1d(python_class, sql_class))
print("only Python (setdiff1d)   :", np.setdiff1d(python_class, sql_class))
print("only one class (setxor1d) :", np.setxor1d(python_class, sql_class))
print("isin: is each Python student also in SQL?", np.isin(python_class, sql_class))

both classes (intersect1d): ['Kiran' 'Neha']
any class (union1d)       : ['Arjun' 'Asha' 'Kiran' 'Meena' 'Neha' 'Ravi']
only Python (setdiff1d)   : ['Asha' 'Ravi']
only one class (setxor1d) : ['Arjun' 'Asha' 'Meena' 'Ravi']
isin: is each Python student also in SQL? [False False  True  True]

The first four return sorted, unique values. isin returns a mask in the original order, so you can filter with it.

9. Real example: product orders

ordered = np.array(["P3", "P1", "P3", "P7", "P1", "P3", "P9"])
catalogue = np.array(["P1", "P2", "P3", "P7"])

products, times = np.unique(ordered, return_counts=True)
best_first = np.argsort(-times)
print("best sellers     :", list(zip(products[best_first].tolist(), times[best_first].tolist())))
print("never ordered    :", np.setdiff1d(catalogue, ordered))
print("not in catalogue :", ordered[~np.isin(ordered, catalogue)])

best sellers     : [('P3', 3), ('P1', 2), ('P7', 1), ('P9', 1)]
never ordered    : ['P2']
not in catalogue : ['P9']

unique with counts plus argsort gives a best-seller list. setdiff1d finds catalogue products nobody ordered, and ~np.isin finds ordered codes that are missing from the catalogue, such as a data-entry mistake.

Exercises

Try these exercises on your own before moving on:

1. Sort np.array([5, 1, 4, 2, 3]) from big to small.
2. Print the names of the 2 lowest-scoring students with argsort.
3. Count how many times each value appears in rng.integers(1, 7, 60).
4. Two lists of customer ids: find customers who bought in January but not in February.
5. Use searchsorted to put ages into bands: child (<13), teen (<20), adult.

Summary

• np.sort returns a sorted copy; arr.sort() sorts in place; reverse with [::-1].
• argsort gives the sorting order, so you can sort related arrays; lexsort sorts by several keys.
• searchsorted finds insert positions in a sorted array, which is ideal for bands.
• partition/argpartition find the top-k without a full sort.
• np.unique gives distinct values, with counts, first positions or codes.
• intersect1d, union1d, setdiff1d, setxor1d and isin compare arrays like sets.

Next: Lesson 13: Matrix Multiplication and np.linalg.