PythonMastery
beginner 18 min read · lesson 3 of 6 in Data Science Fundamentals

NumPy: Working with Numbers

1 · The lesson

read

A Python list of a million numbers is flexible but slow. A NumPy array of a million numbers does maths 10 to 100 times faster, uses much less memory, and looks more like the formula you'd write on paper. The other data science libraries are built on NumPy: pandas, scikit-learn and PyTorch all accept or return NumPy arrays.

This lesson covers the parts you will use most. The full NumPy track goes deeper.


1. Why NumPy Exists

python
import numpy as np
import time

n = 1_000_000

# Pure Python
py = list(range(n))
t0 = time.perf_counter()
py_doubled = [x * 2 for x in py]
py_time = time.perf_counter() - t0

# NumPy
arr = np.arange(n)
t0 = time.perf_counter()
arr_doubled = arr * 2
np_time = time.perf_counter() - t0

print(f"Python list: {py_time*1000:.1f} ms")
print(f"NumPy array: {np_time*1000:.1f} ms")
print(f"Speedup:     {py_time/np_time:.0f}x")

The speed comes from two things. NumPy stores numbers in a contiguous typed buffer (one block of memory, with no Python object for each element). And operations like arr * 2 run in compiled C. You write Python, but the loops run in C.

The convention is import numpy as np. Almost every example online uses it.


2. Creating Arrays

python
import numpy as np

# From a list
a = np.array([1, 2, 3, 4, 5])
print(a)                              # [1 2 3 4 5]

# Constructors
zeros = np.zeros(5)                   # [0. 0. 0. 0. 0.]   floats by default
ones  = np.ones((2, 3))               # 2x3 matrix of 1.0
empty = np.empty(3)                   # fast, but not filled in: values are leftover memory
full  = np.full(4, 7.5)               # [7.5 7.5 7.5 7.5]

# Ranges
r1 = np.arange(0, 10)                 # [0 1 2 3 4 5 6 7 8 9]   like range
r2 = np.arange(0, 1, 0.25)            # [0.  0.25 0.5  0.75]
r3 = np.linspace(0, 1, 5)             # [0.   0.25 0.5  0.75 1. ]  5 points, both ends included

# Random
rng = np.random.default_rng(42)       # random generator with a seed (the modern API)
print(rng.random(3))                  # 3 floats in [0, 1)
print(rng.integers(1, 7, size=10))    # 10 dice rolls

arange excludes the stop value. linspace includes both ends and takes a count of values. Use linspace for plot axes, and arange for indices.


3. dtype, shape, ndim, size

Four attributes describe any array. If you know these four, you can read most NumPy code.

python
a = np.array([[1, 2, 3],
              [4, 5, 6]])

print(a.dtype)                        # int64    every cell has the same type (int32 in this browser)
print(a.shape)                        # (2, 3)   (rows, columns)
print(a.ndim)                         # 2        the number of dimensions
print(a.size)                         # 6        the total number of cells
print(a.itemsize)                     # 8        bytes per cell (4 in this browser)
+ setup added so this can run · re-runs 2 earlier blocks on this page first
# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

np = _AutoMock('np')

# Earlier blocks on this page that this one builds on, re-run quietly first.
try:
    import numpy as np
    import time
except Exception:
    pass
_pm_earlier = [
    'import numpy as np\nimport time\n\nn = 1_000_000\n\n# Pure Python\npy = list(range(n))\nt0 = time.perf_counter()\npy_doubled = [x * 2 for x in py]\npy_time = time.perf_counter() - t0\n\n# NumPy\narr = np.arange(n)\nt0 = time.perf_counter()\narr_doubled = arr * 2\nnp_time = time.perf_counter() - t0\n\nprint(f"Python list: {py_time*1000:.1f} ms")\nprint(f"NumPy array: {np_time*1000:.1f} ms")\nprint(f"Speedup:     {py_time/np_time:.0f}x")',
    'import numpy as np\n\n# From a list\na = np.array([1, 2, 3, 4, 5])\nprint(a)                              # [1 2 3 4 5]\n\n# Constructors\nzeros = np.zeros(5)                   # [0. 0. 0. 0. 0.]   floats by default\nones  = np.ones((2, 3))               # 2x3 matrix of 1.0\nempty = np.empty(3)                   # fast, but not filled in: values are leftover memory\nfull  = np.full(4, 7.5)               # [7.5 7.5 7.5 7.5]\n\n# Ranges\nr1 = np.arange(0, 10)                 # [0 1 2 3 4 5 6 7 8 9]   like range\nr2 = np.arange(0, 1, 0.25)            # [0.  0.25 0.5  0.75]\nr3 = np.linspace(0, 1, 5)             # [0.   0.25 0.5  0.75 1. ]  5 points, both ends included\n\n# Random\nrng = np.random.default_rng(42)       # random generator with a seed (the modern API)\nprint(rng.random(3))                  # 3 floats in [0, 1)\nprint(rng.integers(1, 7, size=10))    # 10 dice rolls',
]
def _pm_run_earlier(_srcs, _g):
    import sys as _s, io as _io, warnings as _w
    for _src in _srcs:
        _o, _e = _s.stdout, _s.stderr
        _s.stdout = _s.stderr = _io.StringIO()
        try:
            with _w.catch_warnings():
                _w.simplefilter('ignore')
                exec(compile(_src, '<earlier block>', 'exec'), _g)
        except (Exception, SystemExit):
            pass
        finally:
            _s.stdout, _s.stderr = _o, _e
    if 'matplotlib.pyplot' in _s.modules:
        _s.modules['matplotlib.pyplot'].close('all')   # their plots belong to their own blocks
_pm_run_earlier(_pm_earlier, globals())
del _pm_earlier, _pm_run_earlier

Common dtypes: int64, float64, bool, int32, float32. Force a dtype when you need to save memory or match an API:

python
b = np.array([1, 2, 3], dtype=np.float32)
print(b.dtype)                        # float32
+ setup added so this can run · re-runs 3 earlier blocks on this page first
# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

np = _AutoMock('np')

# Earlier blocks on this page that this one builds on, re-run quietly first.
try:
    import numpy as np
    import time
except Exception:
    pass
_pm_earlier = [
    'import numpy as np\nimport time\n\nn = 1_000_000\n\n# Pure Python\npy = list(range(n))\nt0 = time.perf_counter()\npy_doubled = [x * 2 for x in py]\npy_time = time.perf_counter() - t0\n\n# NumPy\narr = np.arange(n)\nt0 = time.perf_counter()\narr_doubled = arr * 2\nnp_time = time.perf_counter() - t0\n\nprint(f"Python list: {py_time*1000:.1f} ms")\nprint(f"NumPy array: {np_time*1000:.1f} ms")\nprint(f"Speedup:     {py_time/np_time:.0f}x")',
    'import numpy as np\n\n# From a list\na = np.array([1, 2, 3, 4, 5])\nprint(a)                              # [1 2 3 4 5]\n\n# Constructors\nzeros = np.zeros(5)                   # [0. 0. 0. 0. 0.]   floats by default\nones  = np.ones((2, 3))               # 2x3 matrix of 1.0\nempty = np.empty(3)                   # fast, but not filled in: values are leftover memory\nfull  = np.full(4, 7.5)               # [7.5 7.5 7.5 7.5]\n\n# Ranges\nr1 = np.arange(0, 10)                 # [0 1 2 3 4 5 6 7 8 9]   like range\nr2 = np.arange(0, 1, 0.25)            # [0.  0.25 0.5  0.75]\nr3 = np.linspace(0, 1, 5)             # [0.   0.25 0.5  0.75 1. ]  5 points, both ends included\n\n# Random\nrng = np.random.default_rng(42)       # random generator with a seed (the modern API)\nprint(rng.random(3))                  # 3 floats in [0, 1)\nprint(rng.integers(1, 7, size=10))    # 10 dice rolls',
    'a = np.array([[1, 2, 3],\n              [4, 5, 6]])\n\nprint(a.dtype)                        # int64    every cell has the same type (int32 in this browser)\nprint(a.shape)                        # (2, 3)   (rows, columns)\nprint(a.ndim)                         # 2        the number of dimensions\nprint(a.size)                         # 6        the total number of cells\nprint(a.itemsize)                     # 8        bytes per cell (4 in this browser)',
]
def _pm_run_earlier(_srcs, _g):
    import sys as _s, io as _io, warnings as _w
    for _src in _srcs:
        _o, _e = _s.stdout, _s.stderr
        _s.stdout = _s.stderr = _io.StringIO()
        try:
            with _w.catch_warnings():
                _w.simplefilter('ignore')
                exec(compile(_src, '<earlier block>', 'exec'), _g)
        except (Exception, SystemExit):
            pass
        finally:
            _s.stdout, _s.stderr = _o, _e
    if 'matplotlib.pyplot' in _s.modules:
        _s.modules['matplotlib.pyplot'].close('all')   # their plots belong to their own blocks
_pm_run_earlier(_pm_earlier, globals())
del _pm_earlier, _pm_run_earlier

A list with mixed types turns everything into strings without any warning, and you can't do maths on strings. Watch for it:

python
bad = np.array([1, "two", 3.0])
print(bad.dtype)                      # <U32  an array of strings, not numbers
+ setup added so this can run · re-runs 4 earlier blocks on this page first
# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

np = _AutoMock('np')

# Earlier blocks on this page that this one builds on, re-run quietly first.
try:
    import numpy as np
    import time
except Exception:
    pass
_pm_earlier = [
    'import numpy as np\nimport time\n\nn = 1_000_000\n\n# Pure Python\npy = list(range(n))\nt0 = time.perf_counter()\npy_doubled = [x * 2 for x in py]\npy_time = time.perf_counter() - t0\n\n# NumPy\narr = np.arange(n)\nt0 = time.perf_counter()\narr_doubled = arr * 2\nnp_time = time.perf_counter() - t0\n\nprint(f"Python list: {py_time*1000:.1f} ms")\nprint(f"NumPy array: {np_time*1000:.1f} ms")\nprint(f"Speedup:     {py_time/np_time:.0f}x")',
    'import numpy as np\n\n# From a list\na = np.array([1, 2, 3, 4, 5])\nprint(a)                              # [1 2 3 4 5]\n\n# Constructors\nzeros = np.zeros(5)                   # [0. 0. 0. 0. 0.]   floats by default\nones  = np.ones((2, 3))               # 2x3 matrix of 1.0\nempty = np.empty(3)                   # fast, but not filled in: values are leftover memory\nfull  = np.full(4, 7.5)               # [7.5 7.5 7.5 7.5]\n\n# Ranges\nr1 = np.arange(0, 10)                 # [0 1 2 3 4 5 6 7 8 9]   like range\nr2 = np.arange(0, 1, 0.25)            # [0.  0.25 0.5  0.75]\nr3 = np.linspace(0, 1, 5)             # [0.   0.25 0.5  0.75 1. ]  5 points, both ends included\n\n# Random\nrng = np.random.default_rng(42)       # random generator with a seed (the modern API)\nprint(rng.random(3))                  # 3 floats in [0, 1)\nprint(rng.integers(1, 7, size=10))    # 10 dice rolls',
    'a = np.array([[1, 2, 3],\n              [4, 5, 6]])\n\nprint(a.dtype)                        # int64    every cell has the same type (int32 in this browser)\nprint(a.shape)                        # (2, 3)   (rows, columns)\nprint(a.ndim)                         # 2        the number of dimensions\nprint(a.size)                         # 6        the total number of cells\nprint(a.itemsize)                     # 8        bytes per cell (4 in this browser)',
    'b = np.array([1, 2, 3], dtype=np.float32)\nprint(b.dtype)                        # float32',
]
def _pm_run_earlier(_srcs, _g):
    import sys as _s, io as _io, warnings as _w
    for _src in _srcs:
        _o, _e = _s.stdout, _s.stderr
        _s.stdout = _s.stderr = _io.StringIO()
        try:
            with _w.catch_warnings():
                _w.simplefilter('ignore')
                exec(compile(_src, '<earlier block>', 'exec'), _g)
        except (Exception, SystemExit):
            pass
        finally:
            _s.stdout, _s.stderr = _o, _e
    if 'matplotlib.pyplot' in _s.modules:
        _s.modules['matplotlib.pyplot'].close('all')   # their plots belong to their own blocks
_pm_run_earlier(_pm_earlier, globals())
del _pm_earlier, _pm_run_earlier

4. Vectorised Operations: No For Loops

Every arithmetic operator works element by element. So do the comparison operators and the maths functions.

python
a = np.array([1, 2, 3, 4, 5])
b = np.array([10, 20, 30, 40, 50])

print(a + b)                          # [11 22 33 44 55]
print(a * 2)                          # [ 2  4  6  8 10]
print(b / a)                          # [10. 10. 10. 10. 10.]
print(a ** 2)                         # [ 1  4  9 16 25]

print(np.sqrt(a))                     # [1.    1.414 1.732 2.    2.236]
print(np.log(a))                      # natural log of each
print(np.exp(a))                      # e^x  of each

# Comparisons return boolean arrays
print(a > 2)                          # [False False  True  True  True]
print((a > 1) & (a < 5))              # element-wise AND (use &, not `and`)
+ setup added so this can run · re-runs 5 earlier blocks on this page first
# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

np = _AutoMock('np')

# Earlier blocks on this page that this one builds on, re-run quietly first.
try:
    import numpy as np
    import time
except Exception:
    pass
_pm_earlier = [
    'import numpy as np\nimport time\n\nn = 1_000_000\n\n# Pure Python\npy = list(range(n))\nt0 = time.perf_counter()\npy_doubled = [x * 2 for x in py]\npy_time = time.perf_counter() - t0\n\n# NumPy\narr = np.arange(n)\nt0 = time.perf_counter()\narr_doubled = arr * 2\nnp_time = time.perf_counter() - t0\n\nprint(f"Python list: {py_time*1000:.1f} ms")\nprint(f"NumPy array: {np_time*1000:.1f} ms")\nprint(f"Speedup:     {py_time/np_time:.0f}x")',
    'import numpy as np\n\n# From a list\na = np.array([1, 2, 3, 4, 5])\nprint(a)                              # [1 2 3 4 5]\n\n# Constructors\nzeros = np.zeros(5)                   # [0. 0. 0. 0. 0.]   floats by default\nones  = np.ones((2, 3))               # 2x3 matrix of 1.0\nempty = np.empty(3)                   # fast, but not filled in: values are leftover memory\nfull  = np.full(4, 7.5)               # [7.5 7.5 7.5 7.5]\n\n# Ranges\nr1 = np.arange(0, 10)                 # [0 1 2 3 4 5 6 7 8 9]   like range\nr2 = np.arange(0, 1, 0.25)            # [0.  0.25 0.5  0.75]\nr3 = np.linspace(0, 1, 5)             # [0.   0.25 0.5  0.75 1. ]  5 points, both ends included\n\n# Random\nrng = np.random.default_rng(42)       # random generator with a seed (the modern API)\nprint(rng.random(3))                  # 3 floats in [0, 1)\nprint(rng.integers(1, 7, size=10))    # 10 dice rolls',
    'a = np.array([[1, 2, 3],\n              [4, 5, 6]])\n\nprint(a.dtype)                        # int64    every cell has the same type (int32 in this browser)\nprint(a.shape)                        # (2, 3)   (rows, columns)\nprint(a.ndim)                         # 2        the number of dimensions\nprint(a.size)                         # 6        the total number of cells\nprint(a.itemsize)                     # 8        bytes per cell (4 in this browser)',
    'b = np.array([1, 2, 3], dtype=np.float32)\nprint(b.dtype)                        # float32',
    'bad = np.array([1, "two", 3.0])\nprint(bad.dtype)                      # <U32  an array of strings, not numbers',
]
def _pm_run_earlier(_srcs, _g):
    import sys as _s, io as _io, warnings as _w
    for _src in _srcs:
        _o, _e = _s.stdout, _s.stderr
        _s.stdout = _s.stderr = _io.StringIO()
        try:
            with _w.catch_warnings():
                _w.simplefilter('ignore')
                exec(compile(_src, '<earlier block>', 'exec'), _g)
        except (Exception, SystemExit):
            pass
        finally:
            _s.stdout, _s.stderr = _o, _e
    if 'matplotlib.pyplot' in _s.modules:
        _s.modules['matplotlib.pyplot'].close('all')   # their plots belong to their own blocks
_pm_run_earlier(_pm_earlier, globals())
del _pm_earlier, _pm_run_earlier

The rule: if you are writing a for-loop over a NumPy array, there is probably a better way. Write the operation on the whole array instead.


5. Aggregations

An aggregation reduces an array to one number, or to one number per row or column.

python
a = np.array([[1, 2, 3],
              [4, 5, 6]])

print(a.sum())                        # 21
print(a.mean())                       # 3.5
print(a.max())                        # 6
print(a.min())                        # 1
print(a.std())                        # 1.707...  population standard deviation
print(a.var())                        # 2.916...

# axis=0 collapses the rows, giving one result per column
print(a.sum(axis=0))                  # [5 7 9]
# axis=1 collapses the columns, giving one result per row
print(a.sum(axis=1))                  # [ 6 15]
+ setup added so this can run · re-runs 6 earlier blocks on this page first
# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

np = _AutoMock('np')

# Earlier blocks on this page that this one builds on, re-run quietly first.
try:
    import numpy as np
    import time
except Exception:
    pass
_pm_earlier = [
    'import numpy as np\nimport time\n\nn = 1_000_000\n\n# Pure Python\npy = list(range(n))\nt0 = time.perf_counter()\npy_doubled = [x * 2 for x in py]\npy_time = time.perf_counter() - t0\n\n# NumPy\narr = np.arange(n)\nt0 = time.perf_counter()\narr_doubled = arr * 2\nnp_time = time.perf_counter() - t0\n\nprint(f"Python list: {py_time*1000:.1f} ms")\nprint(f"NumPy array: {np_time*1000:.1f} ms")\nprint(f"Speedup:     {py_time/np_time:.0f}x")',
    'import numpy as np\n\n# From a list\na = np.array([1, 2, 3, 4, 5])\nprint(a)                              # [1 2 3 4 5]\n\n# Constructors\nzeros = np.zeros(5)                   # [0. 0. 0. 0. 0.]   floats by default\nones  = np.ones((2, 3))               # 2x3 matrix of 1.0\nempty = np.empty(3)                   # fast, but not filled in: values are leftover memory\nfull  = np.full(4, 7.5)               # [7.5 7.5 7.5 7.5]\n\n# Ranges\nr1 = np.arange(0, 10)                 # [0 1 2 3 4 5 6 7 8 9]   like range\nr2 = np.arange(0, 1, 0.25)            # [0.  0.25 0.5  0.75]\nr3 = np.linspace(0, 1, 5)             # [0.   0.25 0.5  0.75 1. ]  5 points, both ends included\n\n# Random\nrng = np.random.default_rng(42)       # random generator with a seed (the modern API)\nprint(rng.random(3))                  # 3 floats in [0, 1)\nprint(rng.integers(1, 7, size=10))    # 10 dice rolls',
    'a = np.array([[1, 2, 3],\n              [4, 5, 6]])\n\nprint(a.dtype)                        # int64    every cell has the same type (int32 in this browser)\nprint(a.shape)                        # (2, 3)   (rows, columns)\nprint(a.ndim)                         # 2        the number of dimensions\nprint(a.size)                         # 6        the total number of cells\nprint(a.itemsize)                     # 8        bytes per cell (4 in this browser)',
    'b = np.array([1, 2, 3], dtype=np.float32)\nprint(b.dtype)                        # float32',
    'bad = np.array([1, "two", 3.0])\nprint(bad.dtype)                      # <U32  an array of strings, not numbers',
    'a = np.array([1, 2, 3, 4, 5])\nb = np.array([10, 20, 30, 40, 50])\n\nprint(a + b)                          # [11 22 33 44 55]\nprint(a * 2)                          # [ 2  4  6  8 10]\nprint(b / a)                          # [10. 10. 10. 10. 10.]\nprint(a ** 2)                         # [ 1  4  9 16 25]\n\nprint(np.sqrt(a))                     # [1.    1.414 1.732 2.    2.236]\nprint(np.log(a))                      # natural log of each\nprint(np.exp(a))                      # e^x  of each\n\n# Comparisons return boolean arrays\nprint(a > 2)                          # [False False  True  True  True]\nprint((a > 1) & (a < 5))              # element-wise AND (use &, not `and`)',
]
def _pm_run_earlier(_srcs, _g):
    import sys as _s, io as _io, warnings as _w
    for _src in _srcs:
        _o, _e = _s.stdout, _s.stderr
        _s.stdout = _s.stderr = _io.StringIO()
        try:
            with _w.catch_warnings():
                _w.simplefilter('ignore')
                exec(compile(_src, '<earlier block>', 'exec'), _g)
        except (Exception, SystemExit):
            pass
        finally:
            _s.stdout, _s.stderr = _o, _e
    if 'matplotlib.pyplot' in _s.modules:
        _s.modules['matplotlib.pyplot'].close('all')   # their plots belong to their own blocks
_pm_run_earlier(_pm_earlier, globals())
del _pm_earlier, _pm_run_earlier

axis=0 is "down the rows". axis=1 is "across the columns". The dimension you aggregate over is removed from the output shape.


6. Indexing and Slicing

For 1D arrays, indexing matches Python lists.

python
a = np.arange(10, 60, 10)             # [10 20 30 40 50]
print(a[0])                           # 10
print(a[-1])                          # 50
print(a[1:4])                         # [20 30 40]
print(a[::2])                         # [10 30 50]   every other element
+ setup added so this can run · re-runs 7 earlier blocks on this page first
# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

np = _AutoMock('np')

# Earlier blocks on this page that this one builds on, re-run quietly first.
try:
    import numpy as np
    import time
except Exception:
    pass
_pm_earlier = [
    'import numpy as np\nimport time\n\nn = 1_000_000\n\n# Pure Python\npy = list(range(n))\nt0 = time.perf_counter()\npy_doubled = [x * 2 for x in py]\npy_time = time.perf_counter() - t0\n\n# NumPy\narr = np.arange(n)\nt0 = time.perf_counter()\narr_doubled = arr * 2\nnp_time = time.perf_counter() - t0\n\nprint(f"Python list: {py_time*1000:.1f} ms")\nprint(f"NumPy array: {np_time*1000:.1f} ms")\nprint(f"Speedup:     {py_time/np_time:.0f}x")',
    'import numpy as np\n\n# From a list\na = np.array([1, 2, 3, 4, 5])\nprint(a)                              # [1 2 3 4 5]\n\n# Constructors\nzeros = np.zeros(5)                   # [0. 0. 0. 0. 0.]   floats by default\nones  = np.ones((2, 3))               # 2x3 matrix of 1.0\nempty = np.empty(3)                   # fast, but not filled in: values are leftover memory\nfull  = np.full(4, 7.5)               # [7.5 7.5 7.5 7.5]\n\n# Ranges\nr1 = np.arange(0, 10)                 # [0 1 2 3 4 5 6 7 8 9]   like range\nr2 = np.arange(0, 1, 0.25)            # [0.  0.25 0.5  0.75]\nr3 = np.linspace(0, 1, 5)             # [0.   0.25 0.5  0.75 1. ]  5 points, both ends included\n\n# Random\nrng = np.random.default_rng(42)       # random generator with a seed (the modern API)\nprint(rng.random(3))                  # 3 floats in [0, 1)\nprint(rng.integers(1, 7, size=10))    # 10 dice rolls',
    'a = np.array([[1, 2, 3],\n              [4, 5, 6]])\n\nprint(a.dtype)                        # int64    every cell has the same type (int32 in this browser)\nprint(a.shape)                        # (2, 3)   (rows, columns)\nprint(a.ndim)                         # 2        the number of dimensions\nprint(a.size)                         # 6        the total number of cells\nprint(a.itemsize)                     # 8        bytes per cell (4 in this browser)',
    'b = np.array([1, 2, 3], dtype=np.float32)\nprint(b.dtype)                        # float32',
    'bad = np.array([1, "two", 3.0])\nprint(bad.dtype)                      # <U32  an array of strings, not numbers',
    'a = np.array([1, 2, 3, 4, 5])\nb = np.array([10, 20, 30, 40, 50])\n\nprint(a + b)                          # [11 22 33 44 55]\nprint(a * 2)                          # [ 2  4  6  8 10]\nprint(b / a)                          # [10. 10. 10. 10. 10.]\nprint(a ** 2)                         # [ 1  4  9 16 25]\n\nprint(np.sqrt(a))                     # [1.    1.414 1.732 2.    2.236]\nprint(np.log(a))                      # natural log of each\nprint(np.exp(a))                      # e^x  of each\n\n# Comparisons return boolean arrays\nprint(a > 2)                          # [False False  True  True  True]\nprint((a > 1) & (a < 5))              # element-wise AND (use &, not `and`)',
    'a = np.array([[1, 2, 3],\n              [4, 5, 6]])\n\nprint(a.sum())                        # 21\nprint(a.mean())                       # 3.5\nprint(a.max())                        # 6\nprint(a.min())                        # 1\nprint(a.std())                        # 1.707...  population standard deviation\nprint(a.var())                        # 2.916...\n\n# axis=0 collapses the rows, giving one result per column\nprint(a.sum(axis=0))                  # [5 7 9]\n# axis=1 collapses the columns, giving one result per row\nprint(a.sum(axis=1))                  # [ 6 15]',
]
def _pm_run_earlier(_srcs, _g):
    import sys as _s, io as _io, warnings as _w
    for _src in _srcs:
        _o, _e = _s.stdout, _s.stderr
        _s.stdout = _s.stderr = _io.StringIO()
        try:
            with _w.catch_warnings():
                _w.simplefilter('ignore')
                exec(compile(_src, '<earlier block>', 'exec'), _g)
        except (Exception, SystemExit):
            pass
        finally:
            _s.stdout, _s.stderr = _o, _e
    if 'matplotlib.pyplot' in _s.modules:
        _s.modules['matplotlib.pyplot'].close('all')   # their plots belong to their own blocks
_pm_run_earlier(_pm_earlier, globals())
del _pm_earlier, _pm_run_earlier

For 2D, give two indices: first the row, then the column:

python
m = np.array([[1, 2, 3],
              [4, 5, 6],
              [7, 8, 9]])

print(m[0, 0])                        # 1
print(m[1, 2])                        # 6
print(m[:, 0])                        # [1 4 7]      first column
print(m[1, :])                        # [4 5 6]      second row
print(m[0:2, 1:3])                    # [[2 3]
                                      #  [5 6]]
+ setup added so this can run · re-runs 8 earlier blocks on this page first
# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

np = _AutoMock('np')

# Earlier blocks on this page that this one builds on, re-run quietly first.
try:
    import numpy as np
    import time
except Exception:
    pass
_pm_earlier = [
    'import numpy as np\nimport time\n\nn = 1_000_000\n\n# Pure Python\npy = list(range(n))\nt0 = time.perf_counter()\npy_doubled = [x * 2 for x in py]\npy_time = time.perf_counter() - t0\n\n# NumPy\narr = np.arange(n)\nt0 = time.perf_counter()\narr_doubled = arr * 2\nnp_time = time.perf_counter() - t0\n\nprint(f"Python list: {py_time*1000:.1f} ms")\nprint(f"NumPy array: {np_time*1000:.1f} ms")\nprint(f"Speedup:     {py_time/np_time:.0f}x")',
    'import numpy as np\n\n# From a list\na = np.array([1, 2, 3, 4, 5])\nprint(a)                              # [1 2 3 4 5]\n\n# Constructors\nzeros = np.zeros(5)                   # [0. 0. 0. 0. 0.]   floats by default\nones  = np.ones((2, 3))               # 2x3 matrix of 1.0\nempty = np.empty(3)                   # fast, but not filled in: values are leftover memory\nfull  = np.full(4, 7.5)               # [7.5 7.5 7.5 7.5]\n\n# Ranges\nr1 = np.arange(0, 10)                 # [0 1 2 3 4 5 6 7 8 9]   like range\nr2 = np.arange(0, 1, 0.25)            # [0.  0.25 0.5  0.75]\nr3 = np.linspace(0, 1, 5)             # [0.   0.25 0.5  0.75 1. ]  5 points, both ends included\n\n# Random\nrng = np.random.default_rng(42)       # random generator with a seed (the modern API)\nprint(rng.random(3))                  # 3 floats in [0, 1)\nprint(rng.integers(1, 7, size=10))    # 10 dice rolls',
    'a = np.array([[1, 2, 3],\n              [4, 5, 6]])\n\nprint(a.dtype)                        # int64    every cell has the same type (int32 in this browser)\nprint(a.shape)                        # (2, 3)   (rows, columns)\nprint(a.ndim)                         # 2        the number of dimensions\nprint(a.size)                         # 6        the total number of cells\nprint(a.itemsize)                     # 8        bytes per cell (4 in this browser)',
    'b = np.array([1, 2, 3], dtype=np.float32)\nprint(b.dtype)                        # float32',
    'bad = np.array([1, "two", 3.0])\nprint(bad.dtype)                      # <U32  an array of strings, not numbers',
    'a = np.array([1, 2, 3, 4, 5])\nb = np.array([10, 20, 30, 40, 50])\n\nprint(a + b)                          # [11 22 33 44 55]\nprint(a * 2)                          # [ 2  4  6  8 10]\nprint(b / a)                          # [10. 10. 10. 10. 10.]\nprint(a ** 2)                         # [ 1  4  9 16 25]\n\nprint(np.sqrt(a))                     # [1.    1.414 1.732 2.    2.236]\nprint(np.log(a))                      # natural log of each\nprint(np.exp(a))                      # e^x  of each\n\n# Comparisons return boolean arrays\nprint(a > 2)                          # [False False  True  True  True]\nprint((a > 1) & (a < 5))              # element-wise AND (use &, not `and`)',
    'a = np.array([[1, 2, 3],\n              [4, 5, 6]])\n\nprint(a.sum())                        # 21\nprint(a.mean())                       # 3.5\nprint(a.max())                        # 6\nprint(a.min())                        # 1\nprint(a.std())                        # 1.707...  population standard deviation\nprint(a.var())                        # 2.916...\n\n# axis=0 collapses the rows, giving one result per column\nprint(a.sum(axis=0))                  # [5 7 9]\n# axis=1 collapses the columns, giving one result per row\nprint(a.sum(axis=1))                  # [ 6 15]',
    'a = np.arange(10, 60, 10)             # [10 20 30 40 50]\nprint(a[0])                           # 10\nprint(a[-1])                          # 50\nprint(a[1:4])                         # [20 30 40]\nprint(a[::2])                         # [10 30 50]   every other element',
]
def _pm_run_earlier(_srcs, _g):
    import sys as _s, io as _io, warnings as _w
    for _src in _srcs:
        _o, _e = _s.stdout, _s.stderr
        _s.stdout = _s.stderr = _io.StringIO()
        try:
            with _w.catch_warnings():
                _w.simplefilter('ignore')
                exec(compile(_src, '<earlier block>', 'exec'), _g)
        except (Exception, SystemExit):
            pass
        finally:
            _s.stdout, _s.stderr = _o, _e
    if 'matplotlib.pyplot' in _s.modules:
        _s.modules['matplotlib.pyplot'].close('all')   # their plots belong to their own blocks
_pm_run_earlier(_pm_earlier, globals())
del _pm_earlier, _pm_run_earlier

m[1, :] is the row at index 1. m[:, 0] is the column at index 0. The colon means "all of this axis".


7. Boolean Masking

Boolean masking picks the elements that match a condition. It is one of the most useful ways to index an array.

python
a = np.array([12, 5, 9, 20, 3, 17])

mask = a > 10
print(mask)                           # [ True False False  True False  True]
print(a[mask])                        # [12 20 17]

# The same thing in one line (a common pattern)
print(a[a > 10])                      # [12 20 17]

# Combine conditions with & and | (parentheses required)
print(a[(a > 5) & (a < 18)])          # [12  9 17]

# Assign through a mask: set every value below 5 to 0, in place
a[a < 5] = 0
print(a)                              # [12  5  9 20  0 17]
+ setup added so this can run · re-runs 9 earlier blocks on this page first
# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

np = _AutoMock('np')

# Earlier blocks on this page that this one builds on, re-run quietly first.
try:
    import numpy as np
    import time
except Exception:
    pass
_pm_earlier = [
    'import numpy as np\nimport time\n\nn = 1_000_000\n\n# Pure Python\npy = list(range(n))\nt0 = time.perf_counter()\npy_doubled = [x * 2 for x in py]\npy_time = time.perf_counter() - t0\n\n# NumPy\narr = np.arange(n)\nt0 = time.perf_counter()\narr_doubled = arr * 2\nnp_time = time.perf_counter() - t0\n\nprint(f"Python list: {py_time*1000:.1f} ms")\nprint(f"NumPy array: {np_time*1000:.1f} ms")\nprint(f"Speedup:     {py_time/np_time:.0f}x")',
    'import numpy as np\n\n# From a list\na = np.array([1, 2, 3, 4, 5])\nprint(a)                              # [1 2 3 4 5]\n\n# Constructors\nzeros = np.zeros(5)                   # [0. 0. 0. 0. 0.]   floats by default\nones  = np.ones((2, 3))               # 2x3 matrix of 1.0\nempty = np.empty(3)                   # fast, but not filled in: values are leftover memory\nfull  = np.full(4, 7.5)               # [7.5 7.5 7.5 7.5]\n\n# Ranges\nr1 = np.arange(0, 10)                 # [0 1 2 3 4 5 6 7 8 9]   like range\nr2 = np.arange(0, 1, 0.25)            # [0.  0.25 0.5  0.75]\nr3 = np.linspace(0, 1, 5)             # [0.   0.25 0.5  0.75 1. ]  5 points, both ends included\n\n# Random\nrng = np.random.default_rng(42)       # random generator with a seed (the modern API)\nprint(rng.random(3))                  # 3 floats in [0, 1)\nprint(rng.integers(1, 7, size=10))    # 10 dice rolls',
    'a = np.array([[1, 2, 3],\n              [4, 5, 6]])\n\nprint(a.dtype)                        # int64    every cell has the same type (int32 in this browser)\nprint(a.shape)                        # (2, 3)   (rows, columns)\nprint(a.ndim)                         # 2        the number of dimensions\nprint(a.size)                         # 6        the total number of cells\nprint(a.itemsize)                     # 8        bytes per cell (4 in this browser)',
    'b = np.array([1, 2, 3], dtype=np.float32)\nprint(b.dtype)                        # float32',
    'bad = np.array([1, "two", 3.0])\nprint(bad.dtype)                      # <U32  an array of strings, not numbers',
    'a = np.array([1, 2, 3, 4, 5])\nb = np.array([10, 20, 30, 40, 50])\n\nprint(a + b)                          # [11 22 33 44 55]\nprint(a * 2)                          # [ 2  4  6  8 10]\nprint(b / a)                          # [10. 10. 10. 10. 10.]\nprint(a ** 2)                         # [ 1  4  9 16 25]\n\nprint(np.sqrt(a))                     # [1.    1.414 1.732 2.    2.236]\nprint(np.log(a))                      # natural log of each\nprint(np.exp(a))                      # e^x  of each\n\n# Comparisons return boolean arrays\nprint(a > 2)                          # [False False  True  True  True]\nprint((a > 1) & (a < 5))              # element-wise AND (use &, not `and`)',
    'a = np.array([[1, 2, 3],\n              [4, 5, 6]])\n\nprint(a.sum())                        # 21\nprint(a.mean())                       # 3.5\nprint(a.max())                        # 6\nprint(a.min())                        # 1\nprint(a.std())                        # 1.707...  population standard deviation\nprint(a.var())                        # 2.916...\n\n# axis=0 collapses the rows, giving one result per column\nprint(a.sum(axis=0))                  # [5 7 9]\n# axis=1 collapses the columns, giving one result per row\nprint(a.sum(axis=1))                  # [ 6 15]',
    'a = np.arange(10, 60, 10)             # [10 20 30 40 50]\nprint(a[0])                           # 10\nprint(a[-1])                          # 50\nprint(a[1:4])                         # [20 30 40]\nprint(a[::2])                         # [10 30 50]   every other element',
    'm = np.array([[1, 2, 3],\n              [4, 5, 6],\n              [7, 8, 9]])\n\nprint(m[0, 0])                        # 1\nprint(m[1, 2])                        # 6\nprint(m[:, 0])                        # [1 4 7]      first column\nprint(m[1, :])                        # [4 5 6]      second row\nprint(m[0:2, 1:3])                    # [[2 3]\n                                      #  [5 6]]',
]
def _pm_run_earlier(_srcs, _g):
    import sys as _s, io as _io, warnings as _w
    for _src in _srcs:
        _o, _e = _s.stdout, _s.stderr
        _s.stdout = _s.stderr = _io.StringIO()
        try:
            with _w.catch_warnings():
                _w.simplefilter('ignore')
                exec(compile(_src, '<earlier block>', 'exec'), _g)
        except (Exception, SystemExit):
            pass
        finally:
            _s.stdout, _s.stderr = _o, _e
    if 'matplotlib.pyplot' in _s.modules:
        _s.modules['matplotlib.pyplot'].close('all')   # their plots belong to their own blocks
_pm_run_earlier(_pm_earlier, globals())
del _pm_earlier, _pm_run_earlier

Boolean masking is how you filter a column. In pandas it looks the same: df[df["age"] > 30].


8. Reshaping

Same data, different shape. The total number of elements must match.

python
a = np.arange(12)
print(a.shape)                        # (12,)

b = a.reshape(3, 4)
print(b)
# [[ 0  1  2  3]
#  [ 4  5  6  7]
#  [ 8  9 10 11]]

c = a.reshape(2, 2, 3)                # 3D: 2 layers, each 2x3
print(c.shape)                        # (2, 2, 3)

# -1 means "work out this dimension for me"
d = a.reshape(4, -1)                  # 4 rows, and NumPy works out 3 columns
print(d.shape)                        # (4, 3)
+ setup added so this can run · re-runs 10 earlier blocks on this page first
# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

np = _AutoMock('np')

# Earlier blocks on this page that this one builds on, re-run quietly first.
try:
    import numpy as np
    import time
except Exception:
    pass
_pm_earlier = [
    'import numpy as np\nimport time\n\nn = 1_000_000\n\n# Pure Python\npy = list(range(n))\nt0 = time.perf_counter()\npy_doubled = [x * 2 for x in py]\npy_time = time.perf_counter() - t0\n\n# NumPy\narr = np.arange(n)\nt0 = time.perf_counter()\narr_doubled = arr * 2\nnp_time = time.perf_counter() - t0\n\nprint(f"Python list: {py_time*1000:.1f} ms")\nprint(f"NumPy array: {np_time*1000:.1f} ms")\nprint(f"Speedup:     {py_time/np_time:.0f}x")',
    'import numpy as np\n\n# From a list\na = np.array([1, 2, 3, 4, 5])\nprint(a)                              # [1 2 3 4 5]\n\n# Constructors\nzeros = np.zeros(5)                   # [0. 0. 0. 0. 0.]   floats by default\nones  = np.ones((2, 3))               # 2x3 matrix of 1.0\nempty = np.empty(3)                   # fast, but not filled in: values are leftover memory\nfull  = np.full(4, 7.5)               # [7.5 7.5 7.5 7.5]\n\n# Ranges\nr1 = np.arange(0, 10)                 # [0 1 2 3 4 5 6 7 8 9]   like range\nr2 = np.arange(0, 1, 0.25)            # [0.  0.25 0.5  0.75]\nr3 = np.linspace(0, 1, 5)             # [0.   0.25 0.5  0.75 1. ]  5 points, both ends included\n\n# Random\nrng = np.random.default_rng(42)       # random generator with a seed (the modern API)\nprint(rng.random(3))                  # 3 floats in [0, 1)\nprint(rng.integers(1, 7, size=10))    # 10 dice rolls',
    'a = np.array([[1, 2, 3],\n              [4, 5, 6]])\n\nprint(a.dtype)                        # int64    every cell has the same type (int32 in this browser)\nprint(a.shape)                        # (2, 3)   (rows, columns)\nprint(a.ndim)                         # 2        the number of dimensions\nprint(a.size)                         # 6        the total number of cells\nprint(a.itemsize)                     # 8        bytes per cell (4 in this browser)',
    'b = np.array([1, 2, 3], dtype=np.float32)\nprint(b.dtype)                        # float32',
    'bad = np.array([1, "two", 3.0])\nprint(bad.dtype)                      # <U32  an array of strings, not numbers',
    'a = np.array([1, 2, 3, 4, 5])\nb = np.array([10, 20, 30, 40, 50])\n\nprint(a + b)                          # [11 22 33 44 55]\nprint(a * 2)                          # [ 2  4  6  8 10]\nprint(b / a)                          # [10. 10. 10. 10. 10.]\nprint(a ** 2)                         # [ 1  4  9 16 25]\n\nprint(np.sqrt(a))                     # [1.    1.414 1.732 2.    2.236]\nprint(np.log(a))                      # natural log of each\nprint(np.exp(a))                      # e^x  of each\n\n# Comparisons return boolean arrays\nprint(a > 2)                          # [False False  True  True  True]\nprint((a > 1) & (a < 5))              # element-wise AND (use &, not `and`)',
    'a = np.array([[1, 2, 3],\n              [4, 5, 6]])\n\nprint(a.sum())                        # 21\nprint(a.mean())                       # 3.5\nprint(a.max())                        # 6\nprint(a.min())                        # 1\nprint(a.std())                        # 1.707...  population standard deviation\nprint(a.var())                        # 2.916...\n\n# axis=0 collapses the rows, giving one result per column\nprint(a.sum(axis=0))                  # [5 7 9]\n# axis=1 collapses the columns, giving one result per row\nprint(a.sum(axis=1))                  # [ 6 15]',
    'a = np.arange(10, 60, 10)             # [10 20 30 40 50]\nprint(a[0])                           # 10\nprint(a[-1])                          # 50\nprint(a[1:4])                         # [20 30 40]\nprint(a[::2])                         # [10 30 50]   every other element',
    'm = np.array([[1, 2, 3],\n              [4, 5, 6],\n              [7, 8, 9]])\n\nprint(m[0, 0])                        # 1\nprint(m[1, 2])                        # 6\nprint(m[:, 0])                        # [1 4 7]      first column\nprint(m[1, :])                        # [4 5 6]      second row\nprint(m[0:2, 1:3])                    # [[2 3]\n                                      #  [5 6]]',
    'a = np.array([12, 5, 9, 20, 3, 17])\n\nmask = a > 10\nprint(mask)                           # [ True False False  True False  True]\nprint(a[mask])                        # [12 20 17]\n\n# The same thing in one line (a common pattern)\nprint(a[a > 10])                      # [12 20 17]\n\n# Combine conditions with & and | (parentheses required)\nprint(a[(a > 5) & (a < 18)])          # [12  9 17]\n\n# Assign through a mask: set every value below 5 to 0, in place\na[a < 5] = 0\nprint(a)                              # [12  5  9 20  0 17]',
]
def _pm_run_earlier(_srcs, _g):
    import sys as _s, io as _io, warnings as _w
    for _src in _srcs:
        _o, _e = _s.stdout, _s.stderr
        _s.stdout = _s.stderr = _io.StringIO()
        try:
            with _w.catch_warnings():
                _w.simplefilter('ignore')
                exec(compile(_src, '<earlier block>', 'exec'), _g)
        except (Exception, SystemExit):
            pass
        finally:
            _s.stdout, _s.stderr = _o, _e
    if 'matplotlib.pyplot' in _s.modules:
        _s.modules['matplotlib.pyplot'].close('all')   # their plots belong to their own blocks
_pm_run_earlier(_pm_earlier, globals())
del _pm_earlier, _pm_run_earlier

reshape returns a view when it can, so no data is copied. This makes it very cheap.


9. A Teaser: Broadcasting

What happens when shapes don't match?

python
prices = np.array([[10, 20, 30],
                   [40, 50, 60]])
tax_rate = 0.2

# Scalar + array
print(prices * (1 + tax_rate))
# [[12. 24. 36.]
#  [48. 60. 72.]]

# Row vector + 2D array: the row is added to every row
discount = np.array([0.9, 0.95, 1.0])
print(prices * discount)
# [[ 9.  19.  30. ]
#  [36.  47.5 60. ]]
+ setup added so this can run · re-runs 11 earlier blocks on this page first
# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

np = _AutoMock('np')

# Earlier blocks on this page that this one builds on, re-run quietly first.
try:
    import numpy as np
    import time
except Exception:
    pass
_pm_earlier = [
    'import numpy as np\nimport time\n\nn = 1_000_000\n\n# Pure Python\npy = list(range(n))\nt0 = time.perf_counter()\npy_doubled = [x * 2 for x in py]\npy_time = time.perf_counter() - t0\n\n# NumPy\narr = np.arange(n)\nt0 = time.perf_counter()\narr_doubled = arr * 2\nnp_time = time.perf_counter() - t0\n\nprint(f"Python list: {py_time*1000:.1f} ms")\nprint(f"NumPy array: {np_time*1000:.1f} ms")\nprint(f"Speedup:     {py_time/np_time:.0f}x")',
    'import numpy as np\n\n# From a list\na = np.array([1, 2, 3, 4, 5])\nprint(a)                              # [1 2 3 4 5]\n\n# Constructors\nzeros = np.zeros(5)                   # [0. 0. 0. 0. 0.]   floats by default\nones  = np.ones((2, 3))               # 2x3 matrix of 1.0\nempty = np.empty(3)                   # fast, but not filled in: values are leftover memory\nfull  = np.full(4, 7.5)               # [7.5 7.5 7.5 7.5]\n\n# Ranges\nr1 = np.arange(0, 10)                 # [0 1 2 3 4 5 6 7 8 9]   like range\nr2 = np.arange(0, 1, 0.25)            # [0.  0.25 0.5  0.75]\nr3 = np.linspace(0, 1, 5)             # [0.   0.25 0.5  0.75 1. ]  5 points, both ends included\n\n# Random\nrng = np.random.default_rng(42)       # random generator with a seed (the modern API)\nprint(rng.random(3))                  # 3 floats in [0, 1)\nprint(rng.integers(1, 7, size=10))    # 10 dice rolls',
    'a = np.array([[1, 2, 3],\n              [4, 5, 6]])\n\nprint(a.dtype)                        # int64    every cell has the same type (int32 in this browser)\nprint(a.shape)                        # (2, 3)   (rows, columns)\nprint(a.ndim)                         # 2        the number of dimensions\nprint(a.size)                         # 6        the total number of cells\nprint(a.itemsize)                     # 8        bytes per cell (4 in this browser)',
    'b = np.array([1, 2, 3], dtype=np.float32)\nprint(b.dtype)                        # float32',
    'bad = np.array([1, "two", 3.0])\nprint(bad.dtype)                      # <U32  an array of strings, not numbers',
    'a = np.array([1, 2, 3, 4, 5])\nb = np.array([10, 20, 30, 40, 50])\n\nprint(a + b)                          # [11 22 33 44 55]\nprint(a * 2)                          # [ 2  4  6  8 10]\nprint(b / a)                          # [10. 10. 10. 10. 10.]\nprint(a ** 2)                         # [ 1  4  9 16 25]\n\nprint(np.sqrt(a))                     # [1.    1.414 1.732 2.    2.236]\nprint(np.log(a))                      # natural log of each\nprint(np.exp(a))                      # e^x  of each\n\n# Comparisons return boolean arrays\nprint(a > 2)                          # [False False  True  True  True]\nprint((a > 1) & (a < 5))              # element-wise AND (use &, not `and`)',
    'a = np.array([[1, 2, 3],\n              [4, 5, 6]])\n\nprint(a.sum())                        # 21\nprint(a.mean())                       # 3.5\nprint(a.max())                        # 6\nprint(a.min())                        # 1\nprint(a.std())                        # 1.707...  population standard deviation\nprint(a.var())                        # 2.916...\n\n# axis=0 collapses the rows, giving one result per column\nprint(a.sum(axis=0))                  # [5 7 9]\n# axis=1 collapses the columns, giving one result per row\nprint(a.sum(axis=1))                  # [ 6 15]',
    'a = np.arange(10, 60, 10)             # [10 20 30 40 50]\nprint(a[0])                           # 10\nprint(a[-1])                          # 50\nprint(a[1:4])                         # [20 30 40]\nprint(a[::2])                         # [10 30 50]   every other element',
    'm = np.array([[1, 2, 3],\n              [4, 5, 6],\n              [7, 8, 9]])\n\nprint(m[0, 0])                        # 1\nprint(m[1, 2])                        # 6\nprint(m[:, 0])                        # [1 4 7]      first column\nprint(m[1, :])                        # [4 5 6]      second row\nprint(m[0:2, 1:3])                    # [[2 3]\n                                      #  [5 6]]',
    'a = np.array([12, 5, 9, 20, 3, 17])\n\nmask = a > 10\nprint(mask)                           # [ True False False  True False  True]\nprint(a[mask])                        # [12 20 17]\n\n# The same thing in one line (a common pattern)\nprint(a[a > 10])                      # [12 20 17]\n\n# Combine conditions with & and | (parentheses required)\nprint(a[(a > 5) & (a < 18)])          # [12  9 17]\n\n# Assign through a mask: set every value below 5 to 0, in place\na[a < 5] = 0\nprint(a)                              # [12  5  9 20  0 17]',
    'a = np.arange(12)\nprint(a.shape)                        # (12,)\n\nb = a.reshape(3, 4)\nprint(b)\n# [[ 0  1  2  3]\n#  [ 4  5  6  7]\n#  [ 8  9 10 11]]\n\nc = a.reshape(2, 2, 3)                # 3D: 2 layers, each 2x3\nprint(c.shape)                        # (2, 2, 3)\n\n# -1 means "work out this dimension for me"\nd = a.reshape(4, -1)                  # 4 rows, and NumPy works out 3 columns\nprint(d.shape)                        # (4, 3)',
]
def _pm_run_earlier(_srcs, _g):
    import sys as _s, io as _io, warnings as _w
    for _src in _srcs:
        _o, _e = _s.stdout, _s.stderr
        _s.stdout = _s.stderr = _io.StringIO()
        try:
            with _w.catch_warnings():
                _w.simplefilter('ignore')
                exec(compile(_src, '<earlier block>', 'exec'), _g)
        except (Exception, SystemExit):
            pass
        finally:
            _s.stdout, _s.stderr = _o, _e
    if 'matplotlib.pyplot' in _s.modules:
        _s.modules['matplotlib.pyplot'].close('all')   # their plots belong to their own blocks
_pm_run_earlier(_pm_earlier, globals())
del _pm_earlier, _pm_run_earlier

NumPy stretches the smaller shape to match the larger one without copying. This is called broadcasting, and it lets you write short array code without loops. The advanced NumPy lesson covers it in full.


Common Mistakes

  • Pythonic for-loops over arrays. for i in range(len(arr)): arr[i] = arr[i] * 2 instead of arr * 2. This loses all of NumPy's speed advantage.
  • Mixed-type lists become string arrays. np.array([1, "two", 3]) gives you an array of strings. Convert the types before you create the array.
  • Two kinds of division on integer arrays. np.array([1, 2, 3]) / 2 returns floats. np.array([1, 2, 3]) // 2 is floor division: it rounds down to whole numbers. Know which one you want.
  • Using and / or on arrays. They raise ValueError. Use & and |, and put each condition in parentheses.
  • Forgetting axis=. a.sum() collapses the whole array. a.sum(axis=0) returns one value per column.
  • Confusing np.arange with np.linspace. arange takes a step, linspace takes a count. Use linspace when you care about the endpoints.

🎯 Your Turn: normalize()

Implement z-score normalisation: for each value subtract the mean and divide by the standard deviation. A normalised array has mean 0 and standard deviation 1. You will do this preprocessing step very often before ML.

Skeleton:

python
import numpy as np

def normalize(arr):
    """Return (arr - mean) / std. Handle std == 0 sensibly."""
    arr = np.asarray(arr, dtype=float)
    # TODO 1: compute mean and std
    # TODO 2: subtract and divide on the whole array, with no loops
    ...

# Test
data = np.array([10, 20, 30, 40, 50])
z = normalize(data)
print(z)                              # roughly [-1.41 -0.71  0.    0.71  1.41]
print(z.mean(), z.std())              # ~0.0, 1.0
Hint 1: Two numbers, then one expression Compute m = arr.mean() and s = arr.std() first. Then return (arr - m) / s. NumPy broadcasts the scalars across the whole array, so no loop is needed.
Hint 2: Constant arrays If every element is the same, std is 0 and you'd divide by zero. Decide what to do: return zeros, raise ValueError, or replace std with 1. Returning zeros is the most useful default, because it doesn't crash the caller.
Show full solution
python
import numpy as np

def normalize(arr):
    arr = np.asarray(arr, dtype=float)
    mean = arr.mean()
    std  = arr.std()
    if std == 0:
        return np.zeros_like(arr)     # all values are the same, so there is no spread
    return (arr - mean) / std

data = np.array([10, 20, 30, 40, 50])
z = normalize(data)
print(z)                              # [-1.41 -0.71  0.    0.71  1.41]
print(round(z.mean(), 10), round(z.std(), 10))   # 0.0 1.0

# Edge case
print(normalize([7, 7, 7]))           # [0. 0. 0.]

The real work is three lines: np.asarray, two scalar reductions, and one broadcast subtract-and-divide. Scikit-learn's StandardScaler is the same calculation with a fit/transform interface around it.


What You Learned

  • NumPy arrays are contiguous, typed and fast: 10-100x quicker than Python lists for numeric work.
  • Four attributes describe an array: dtype, shape, ndim, size.
  • Vectorised operations replace for-loops: a + b, a * 2, np.sqrt(a).
  • Aggregations collapse arrays, and axis= controls which dimension disappears.
  • Boolean masking filters by condition: a[a > 5].
  • Reshape changes the shape of the data, usually without copying it.
  • Broadcasting lets arrays of different shapes work together, so you can write short array code without loops.

Next: Pandas: Working with Tables: NumPy with column names, indexes, and the file reading and writing you need.

Practice this

on practicepython.in

Short exercises that run in your browser and tell you what your code did, as well as whether a test passed.

Found this useful? Share it with a friend who’s learning Python.