Probability Basics for Data Analysis (With Python Simulations)

Probability measures how likely something is, on a scale from 0 (impossible) to 1 (certain). It underpins everything that follows in statistics: hypothesis tests, A/B tests and machine learning all reason about chance.

A great way to learn probability is to calculate an answer and then check it by simulation: repeat the random experiment many times with NumPy and count.

Basic probability#

When all outcomes are equally likely:

P(event) = (number of favourable outcomes) ÷ (total number of outcomes)

For a fair die, P(rolling an even number) = 3/6 = 0.5. Simulating 100,000 rolls agrees:

import numpy as np

rng = np.random.default_rng(10)
rolls = rng.integers(1, 7, size=100_000)
print(round((rolls % 2 == 0).mean(), 3))
print(round((rolls >= 5).mean(), 3), round(2 / 6, 3))

0.5
0.333 0.333

Combining events#

• Not: P(not A) = 1 − P(A).
• Or (addition rule): P(A or B) = P(A) + P(B) − P(A and B). Subtract the overlap so it is not counted twice.
• And for independent events (one does not affect the other): P(A and B) = P(A) × P(B).

The chance of at least one six in four rolls is easiest through "not": 1 − (5/6)⁴.

import numpy as np

rng = np.random.default_rng(10)
exact = 1 - (5 / 6) ** 4
four_rolls = rng.integers(1, 7, size=(100_000, 4))
simulated = (four_rolls == 6).any(axis=1).mean()
print(round(exact, 4), round(simulated, 4))

0.5177 0.5173

Conditional probability#

P(A | B), "the probability of A given B", restricts attention to the cases where B happened:

P(A | B) = P(A and B) ÷ P(B)

With real data, a conditional probability is just a proportion within a filtered group:

import pandas as pd

visits = pd.DataFrame({
    "device": ["mobile"] * 600 + ["desktop"] * 400,
    "bought": [True] * 30 + [False] * 570 + [True] * 40 + [False] * 360,
})
print(round(visits["bought"].mean(), 3))                                   # P(buy)
print(visits.groupby("device")["bought"].mean().round(3))                  # P(buy | device)
print(round((visits["device"] == "mobile")[visits["bought"]].mean(), 3))   # P(mobile | buy)

0.07
device
desktop    0.10
mobile     0.05
Name: bought, dtype: float64
0.429

Desktop visitors convert at twice the rate of mobile visitors (10% vs 5%). And although 60% of visitors use mobile, only 43% of buyers do (30 of the 70). P(buy | mobile) and P(mobile | buy) are different questions with different answers.

Bayes' theorem#

Bayes' theorem turns one conditional probability around:

P(A | B) = P(B | A) × P(A) ÷ P(B)

Classic example: a disease affects 1% of people. A test catches 95% of real cases but also flags 5% of healthy people. If you test positive, how likely are you to have the disease?

import numpy as np

p_disease = 0.01
p_pos_given_disease = 0.95
p_pos_given_healthy = 0.05

p_pos = p_pos_given_disease * p_disease + p_pos_given_healthy * (1 - p_disease)
print(round(p_pos_given_disease * p_disease / p_pos, 3))

rng = np.random.default_rng(5)
n = 1_000_000
sick = rng.random(n) < p_disease
positive = np.where(sick, rng.random(n) < 0.95, rng.random(n) < 0.05)
print(round(sick[positive].mean(), 3))

0.161
0.162

Only about 16%! Because the disease is rare, most positive results come from the large healthy group. Ignoring how common something is in the first place (the base rate) is one of the most common reasoning errors, in medicine, fraud detection and spam filtering alike.

Expected value#

The expected value is the long-run average outcome: each outcome times its probability, added up. A raffle ticket costs ₹100; there is a 1-in-500 chance to win ₹20,000 and a 1-in-50 chance to win ₹1,000:

import numpy as np

outcomes = np.array([20_000, 1_000, 0])
probs = np.array([1 / 500, 1 / 50, 1 - 1 / 500 - 1 / 50])
expected_prize = (outcomes * probs).sum()
print(round(expected_prize, 2), "expected prize per ticket")
print(round(expected_prize - 100, 2), "expected profit per ticket")

60.0 expected prize per ticket
-40.0 expected profit per ticket

On average you lose ₹40 per ticket. Businesses use the same calculation for pricing, insurance and deciding which experiments to run.

The binomial distribution#

When you repeat a yes/no trial n times with success probability p (independent each time), the number of successes follows a binomial distribution. If 4% of visitors click an ad, what about 50 visitors?

from scipy.stats import binom

n, p = 50, 0.04
print(round(binom.pmf(0, n, p), 3))          # exactly 0 clicks
print(round(binom.pmf(2, n, p), 3))          # exactly 2 clicks
print(round(1 - binom.cdf(4, n, p), 3))      # 5 or more clicks
print(binom.mean(n, p))                      # expected clicks

0.13
0.276
0.049
2.0

pmf gives the probability of an exact count; cdf the probability of at most that count.

Common mistakes#

Practice#

import numpy as np

all_different = np.prod((365 - np.arange(30)) / 365)
print(round(1 - all_different, 3))

rng = np.random.default_rng(1)
classes = rng.integers(0, 365, size=(20_000, 30))
shared = np.array([len(set(c)) < 30 for c in classes])
print(round(shared.mean(), 3))

0.706
0.707

Key takeaway#

Next: deciding whether a difference is real, with Hypothesis Testing.