Skip to main content

NumPy: Vectorized Operations, Reshaping & Linear Algebra

Introduction​

The previous chapter covered creating and indexing arrays; this one covers operating on them. NumPy operations are vectorized — they apply to every element automatically, without a Python for loop, which is both faster and more readable. This chapter also covers broadcasting (operating on arrays of different but compatible shapes), reshaping arrays for different views of the same data, and a few linear-algebra operations that come up constantly in financial/business calculations — dot products, correlation, sorting and ranking.

Concepts​

Element-wise arithmetic​

  • Arithmetic between two same-shaped arrays (a * b, a - b) operates element-by-element.
  • Arithmetic between an array and a scalar (prices * 1.1) applies the scalar to every element.

Broadcasting​

  • NumPy can operate on arrays of different shapes as long as they're "compatible" — dimensions match, or one of them is size 1.
  • The classic pattern: a (4, 7) matrix and a (4, 1) column vector — reshape the vector to (4, 1) with .reshape(4, 1) so each row gets its own scalar applied across every column.

Aggregate functions​

  • .sum(), .mean(), .std(), .min(), .max(), np.percentile(), np.median() — most accept an axis parameter: axis=0 aggregates down columns, axis=1 aggregates across rows.
  • .argmax() / .argmin() return the index of the max/min value, not the value itself.

Reshaping, stacking, splitting​

  • .reshape(new_shape) changes an array's shape without changing its data; -1 in one dimension means "infer this from the total size."
  • .T transposes a 2-D array (rows become columns).
  • np.vstack / np.hstack stack arrays vertically (add rows) or horizontally (add columns); np.split divides one array into several.

Linear algebra and ranking​

  • @ (or np.dot) performs matrix multiplication — for a (rows, cols) units matrix and a (cols,) price vector, units @ prices gives each row's total revenue in one operation.
  • np.corrcoef(a, b) computes the Pearson correlation matrix between two (or more) series.
  • np.argsort(arr) returns the indices that would sort the array — apply those indices to a second, related array (e.g. labels) to sort it in the same order.

Examples​

Element-wise arithmetic:

import numpy as np

units  = np.array([120, 95, 142, 88, 210, 175, 133])
prices = np.array([499, 799, 299, 999, 199, 599, 449])

revenue = units * prices                       # element-wise multiplication
discounted_prices = prices * (1 - 0.10)        # scalar broadcast to every element
cost = units * prices * 0.40
gross_profit = revenue - cost
print(gross_profit.sum())                      # total week's gross profit

Broadcasting a per-store value across a matrix:

weekly_units = np.array([
    [120, 95, 142, 88, 210, 175, 133],
    [98,  112, 89, 104, 188, 160, 120],
    [75,  80,  95, 70,  155, 130, 110],
    [50,  60,  55, 65,  120, 95,  80],
])
store_prices = np.array([499, 799, 299, 999])          # shape (4,)

# Reshape to (4, 1) so it broadcasts down each store's row, across all 7 days
revenue_matrix = weekly_units * store_prices.reshape(4, 1)
print(revenue_matrix.shape)   # (4, 7)

discounts = np.array([0.10, 0.15, 0.05, 0.20]).reshape(4, 1)
net_revenue = revenue_matrix * (1 - discounts)

Aggregates along an axis:

store_totals = weekly_units.sum(axis=1)   # total per store (across the 7 days)
daily_totals = weekly_units.sum(axis=0)   # total per day (across the 4 stores)

best_store_idx = store_totals.argmax()
stores = ['Mumbai', 'Delhi', 'Pune', 'Jaipur']
print('Best store:', stores[best_store_idx])

Reshaping and stacking:

flat = weekly_units.reshape(-1)        # (28,) — flattened for a statistical test

monthly_rev = np.arange(100, 1300, 100)
quarterly = monthly_rev.reshape(4, 3)  # 4 quarters x 3 months
print(quarterly.T)                     # transposed: shape (3, 4)

new_stores = np.array([
    [65, 72, 80, 60, 140, 115, 95],   # Nagpur
    [45, 55, 60, 50, 100,  88, 75],   # Surat
])
all_stores = np.vstack([weekly_units, new_stores])   # (6, 7)

Matrix multiplication and correlation:

units_matrix = np.array([
    [10, 5, 8, 3, 12],
    [8, 7, 6, 5, 10],
    [6, 4, 9, 2, 8],
    [5, 3, 7, 4, 6],
])
prices_vector = np.array([499, 999, 299, 1499, 199])

store_revenue = units_matrix @ prices_vector   # each store's total revenue
corr = np.corrcoef(weekly_units)               # how correlated are stores' sales?

Sorting and ranking:

store_annual = np.array([4820000, 3950000, 5610000, 2890000, 4100000, 3750000])
store_names  = np.array(['Mumbai', 'Delhi', 'Pune', 'Jaipur', 'Nagpur', 'Surat'])

sort_idx = np.argsort(store_annual)
print('Ranked (low to high):', store_names[sort_idx])

top3_idx = np.argsort(store_annual)[-3:][::-1]
print('Top 3:', store_names[top3_idx])

Challenges​

1. Array operations​

Using this revenue data for 3 product categories over 4 quarters:

revenue = np.array([[120000, 145000, 130000, 180000],   # Electronics
                    [ 85000,  92000,  78000, 110000],   # Clothing
                    [ 60000,  75000,  68000,  95000]])  # Home & Kitchen
  • Calculate total annual revenue per category.
  • Calculate quarterly revenue totals across all categories.
  • Find which category-quarter combination had maximum revenue (argmax with unravel_index).
  • Compute quarter-over-quarter growth rate for Electronics.
  • Normalize the matrix so each cell = its value / total revenue (percentage share).
💡 Show Answer
import numpy as np

revenue = np.array([[120000, 145000, 130000, 180000],
                     [ 85000,  92000,  78000, 110000],
                     [ 60000,  75000,  68000,  95000]])

annual_per_category = revenue.sum(axis=1)
print('Annual per category:', annual_per_category)

quarterly_totals = revenue.sum(axis=0)
print('Quarterly totals:', quarterly_totals)

flat_max_idx = revenue.argmax()
category_idx, quarter_idx = np.unravel_index(flat_max_idx, revenue.shape)
print(f'Max revenue: category {category_idx}, quarter {quarter_idx}')

electronics = revenue[0]
qoq_growth = np.diff(electronics) / electronics[:-1] * 100
print('Electronics QoQ growth %:', qoq_growth.round(2))

normalized = revenue / revenue.sum()
print('Percentage share:', (normalized * 100).round(2))

2. Advanced NumPy capstone​

RetailEdge wants a full performance matrix. Given units_sold and unit_price, both shape (6, 12) (6 stores x 12 months):

  1. Compute revenue_matrix = units_sold * unit_price.
  2. Apply store-level discounts using broadcasting: discounts shape (6, 1).
  3. Compute net_revenue and rank stores by annual net revenue.
  4. Find which month had the highest total revenue across all stores.
  5. Compute the Pearson correlation between two stores' monthly revenues.
  6. Simulate 10,000 random daily sales and compute the 5th and 95th percentiles.
💡 Show Answer
import numpy as np

np.random.seed(7)
units_sold = np.random.randint(50, 300, size=(6, 12))
unit_price = np.random.uniform(200, 900, size=(6, 12))
store_names = np.array(['Mumbai', 'Delhi', 'Pune', 'Jaipur', 'Nagpur', 'Surat'])

# 1. Revenue
revenue_matrix = units_sold * unit_price

# 2. Apply per-store discounts via broadcasting
discounts = np.random.uniform(0.05, 0.20, size=(6, 1))
net_revenue = revenue_matrix * (1 - discounts)

# 3. Rank stores by annual net revenue
annual_net = net_revenue.sum(axis=1)
rank_idx = np.argsort(annual_net)[::-1]
print('Store ranking (highest first):', store_names[rank_idx])

# 4. Best month across all stores
monthly_totals = net_revenue.sum(axis=0)
print('Best month index:', monthly_totals.argmax())

# 5. Correlation between two stores (Mumbai=row 0, Delhi=row 1)
corr = np.corrcoef(net_revenue[0], net_revenue[1])[0, 1]
print('Mumbai-Delhi correlation:', round(corr, 3))

# 6. Simulated daily sales percentiles
simulated_daily = np.random.uniform(5000, 50000, size=10000)
p5, p95 = np.percentile(simulated_daily, [5, 95])
print(f'5th percentile: {p5:.0f}, 95th percentile: {p95:.0f}')