Python NumPy Basics: Arrays, Indexing and Calculations

Python with AI Tutorial · Chapter 37 of 48

Python NumPy is the library for fast numerical work. Its core object, the ndarray, stores numbers of one type in a compact block of memory and lets you calculate on thousands of values with a single expression instead of a loop. NumPy also powers pandas, Matplotlib and most data science tools, so learning it pays off everywhere.

Install NumPy and create an array

Install once with pip, then import it as np, the convention every tutorial follows. np.array() turns a Python list into an array.

pip install numpy

Output:
Successfully installed numpy-2.1.1

import numpy as np

units = np.array([4, 10, 2, 15, 7, 9])
print(units)
print(type(units).__name__, units.dtype, units.shape, units.ndim)

Output:
[ 4 10  2 15  7  9]
ndarray int64 (6,) 1

Every array has a dtype (all elements share it), a shape (a tuple of sizes per dimension) and ndim (the number of dimensions). This one is a 1-D array of six 64-bit integers. Lists can hold anything, but arrays trade that flexibility for speed: a million-element array sums in about a millisecond.

Vectorised calculations

Arithmetic applies to every element at once. Multiplying two arrays multiplies them pairwise, and the aggregate methods return totals without a loop.

prices = np.array([999, 299, 999, 99, 299, 149])
revenue = units * prices
print(revenue)
print(revenue.sum(), revenue.mean().round(2), revenue.max())
print(np.round(revenue * 1.18, 2))      # add 18% tax

Output:
[3996 2990 1998 1485 2093 1341]
13903 2317.17 3996
[4715.28 3528.2  2357.64 1752.3  2469.74 1582.38]

Multiplying by the scalar 1.18 is called broadcasting: NumPy stretches the single number across the whole array. The result is a new float array; the original integers are unchanged.

Indexing and slicing

Arrays index like lists, with two additions: a list of positions selects several elements, and a Boolean array selects the elements where it is True.

print(units[0], units[-1])
print(units[1:4])
print(units[[0, 2, 4]])
print(revenue[revenue > 2000])

Output:
4 9
[10  2 15]
[4 2 7]
[3996 2990 2093]

Boolean masks with business rules

A comparison produces a mask. Masks can be combined with & (and), | (or) and ~ (not), and np.where() picks one of two values per element.

regions = np.array(["North", "South", "East", "North", "West", "South"])
north = regions == "North"
print(north)
print("North revenue:", revenue[north].sum())
print(revenue[(regions == "South") & (revenue > 2000)])
print(np.where(revenue > 2000, "High", "Low"))

Output:
[ True False False  True False False]
North revenue: 5481
[2990]
['High' 'High' 'Low' 'Low' 'High' 'Low']

Two-dimensional arrays

A list of lists becomes a 2-D array, like a sheet with rows and columns. Index with [row, column], and use : to take a whole row or column. axis=0 aggregates down the rows, axis=1 across the columns.

sales = np.array([[4, 999],
                  [10, 299],
                  [2, 999],
                  [15, 99]])       # columns: Units, Price
print(sales.shape)
print(sales[:, 0])                 # all rows, first column
print(sales[1])                    # second row
print((sales[:, 0] * sales[:, 1]).sum())
print(sales.sum(axis=0))

Output:
(4, 2)
[ 4 10  2 15]
[ 10 299]
10469
[  31 2396]

Read axis=0 as “collapse the rows”: the result has one number per column, here total units and total price. axis=1 would give one number per row instead.

Here is my Python code that calculates commission for each sales rep:

revenue = [3996, 2990, 1998, 1485, 2093, 1341]
commission = []
for r in revenue:
    if r > 3000:
        commission.append(r * 0.10)
    elif r > 2000:
        commission.append(r * 0.07)
    else:
        commission.append(r * 0.05)

Rewrite it with NumPy so there is no Python loop (use np.select or np.where), print the result rounded to 2 decimals, and explain in three sentences why the NumPy version is faster on 1 million rows.

Creating arrays quickly

You rarely type every value. arange, linspace, zeros and ones build arrays of any size, and reshape changes the shape without copying data.

months = np.arange(1, 13)
print(months.reshape(4, 3))
targets = np.linspace(1000, 2000, 5)
print(targets)
print(np.zeros(3), np.ones((2, 2)))

Output:
[[ 1  2  3]
 [ 4  5  6]
 [ 7  8  9]
 [10 11 12]]
[1000. 1250. 1500. 1750. 2000.]
[0. 0. 0.] [[1. 1.]
 [1. 1.]]

arange() works like range() but returns an array and accepts decimal steps. linspace() is the better choice when you know how many points you need rather than the step size. The new shape passed to reshape() must hold exactly the same number of elements, otherwise NumPy raises a ValueError.

NumPy data types

Because an array has one dtype, NumPy can store it compactly and run arithmetic in compiled C. Check .dtype when results look wrong.

Mixing numbers and text in one np.array() silently converts everything to strings. Keep numbers and labels in separate arrays, or move to a pandas DataFrame, which allows one dtype per column.

Statistics in one line

NumPy has functions for every common statistic. They accept an axis argument on 2-D arrays as well.

print(np.median(revenue), np.std(revenue).round(2))
print(np.percentile(revenue, 75))
print(np.cumsum(revenue))
best = np.argmax(revenue)
print(best, regions[best], revenue[best])

Output:
2045.5 919.12
2765.75
[ 3996  6986  8984 10469 12562 13903]
0 North 3996

np.std() is the population standard deviation. For the sample version that Excel’s STDEV.S returns, pass ddof=1. argmax returns the position of the largest value, which you can then use to look up the matching label.

I have a NumPy array sales with shape (4, 2) holding Units and Price per row, and an array discounts with shape (4,) holding a percentage discount per row. sales * discounts works, but sales * np.array([0.9, 0.8]) also works and gives a different result. Explain NumPy broadcasting rules with these exact shapes, show which axis each array is stretched along, and tell me how to reshape discounts so it applies per row versus per column.

Common mistakes

• Using and / or instead of & / | to combine masks.
• Expecting arr.mean() to skip missing values. NumPy propagates nan; use np.nanmean() or clean the data first.
• Assuming a slice is a copy. view = units[1:4]; view[0] = 0 changes units too. Call .copy() when you need independence.
• Storing text and numbers in one array, which turns every number into a string.
• Looping over elements with for when a vectorised expression or np.where() does the same job many times faster.

Exercise

Using the units, prices and regions arrays from this chapter, calculate revenue, then print the regions whose revenue is above the mean and the percentage of total revenue they contribute, rounded to one decimal.

import numpy as np

units = np.array([4, 10, 2, 15, 7, 9])
prices = np.array([999, 299, 999, 99, 299, 149])
regions = np.array(["North", "South", "East", "North", "West", "South"])

revenue = units * prices
above = revenue > revenue.mean()
print(regions[above])
print(round(revenue[above].sum() / revenue.sum() * 100, 1), "%")

Output:
['North' 'South']
50.2 %

Related chapters

• Python Lists – the built-in structure NumPy arrays replace for numeric work.
• Python pandas – DataFrames built on top of NumPy.
• Python Math – the standard library functions NumPy generalises to arrays.
• Python with AI course hub – all 48 chapters in order.

FAQ

What is NumPy used for in Python?

NumPy provides the ndarray, a fast fixed-type array, plus vectorised maths, statistics, linear algebra and random numbers. It is the numerical foundation for pandas, Matplotlib, scikit-learn and most scientific Python libraries.

What is the difference between a NumPy array and a Python list?

A list can hold mixed types and grows dynamically, but arithmetic needs a loop. A NumPy array holds one dtype in contiguous memory, so operations run element-wise in compiled code, typically 10 to 100 times faster and using far less memory.

Do I need to learn NumPy before pandas?

A basic understanding helps. pandas columns are NumPy arrays underneath, so concepts such as dtype, Boolean masks, axis and vectorised operations carry straight over to DataFrames.

Working with spreadsheets too? Ready-made Excel, Google Sheets and Power BI templates are at NextGenTemplates.com.

Chapter 37 of 48 · Python with AI: all 48 chapters