Segmentation Forests

Unsupervised segment discovery using divergence-based decision trees inspired by Random Forests.

Automatically discover meaningful segments in your data where metric distributions significantly diverge from the background. Perfect for exploratory data analysis, anomaly detection, and customer segmentation without requiring labeled data.

🎯 The Problem

You have a dataset with many categorical features and a metric you care about. For example:

E-commerce: Users across countries, devices, age groups → conversion rate
Digital advertising: Impressions across demographics, platforms, times → CTR
Healthcare: Patients across conditions, treatments, demographics → readmission rate
Finance: Transactions across customer segments, times → fraud rate

The question: Which specific combinations of features exhibit unusual behavior?

With N features and many values per feature, exhaustively testing combinations is impossible. Segmentation Forests solves this by intelligently searching for segments where your metric's distribution differs significantly from the overall population.

✨ Key Features

🌳 Tree-based discovery: Greedy algorithm efficiently navigates combinatorial feature space
🌲 Forest ensemble: Bootstrap aggregating for robust, reproducible patterns
📊 Statistical rigor: KS distance (continuous) & Jensen-Shannon divergence (discrete)
📈 Beautiful visualizations: Distribution comparisons and quality assessments
🔬 Fully typed: Complete type hints for excellent IDE support
⚡ Fast & scalable: Handles datasets with millions of rows

🚀 Quick Start

Installation

pip install segmentation-forests

Basic Example

import pandas as pd
from segmentation_forests import SegmentationTree, SegmentationForest

# Your data: features + metric
data = pd.DataFrame({
    'country': ['US', 'UK', 'US', 'UK', ...],
    'device': ['Mobile', 'Desktop', 'Mobile', ...],
    'gender': ['F', 'M', 'F', ...],
    'impressions': [245, 103, 312, 98, ...]  # Your metric
})

# Discover segments with a single tree
tree = SegmentationTree(max_depth=3, min_samples_split=100)
tree.fit(data, metric_column='impressions')
segments = tree.get_segments(min_divergence=0.1)

# View results
for i, seg in enumerate(segments[:3], 1):
    print(f"{i}. {seg.get_condition_string()}")
    print(f"   Divergence: {seg.divergence:.3f} | Size: {seg.size:,}")

Output:

1. gender == F AND device == Mobile AND country == UK
   Divergence: 0.948 | Size: 523

2. time_of_day == Evening AND country == US AND device == Desktop
   Divergence: 0.856 | Size: 412

3. country == DE AND time_of_day == Morning
   Divergence: 0.824 | Size: 289

🌲 Using the Forest (Recommended)

For more robust results, use the ensemble approach:

from segmentation_forests import SegmentationForest

# Create forest with bootstrap sampling and random features
forest = SegmentationForest(
    n_trees=10,
    max_depth=3,
    max_features=2,  # Random feature selection
    min_samples_split=100,
    min_samples_leaf=50
)

# Fit and get segments found by multiple trees
forest.fit(data, metric_column='impressions')
robust_segments = forest.get_segments(min_support=3, min_divergence=0.1)

# View results
for seg in robust_segments:
    cond_str = " AND ".join([f"{c[0]} {c[1]} {c[2]}" for c in seg['conditions']])
    print(f"{cond_str}")
    print(f"  Support: {seg['support']}/10 trees ({seg['support_rate']*100:.0f}%)")
    print(f"  Avg Divergence: {seg['avg_divergence']:.3f}")
    print()

📊 Visualization

Beautiful distribution comparison plots:

from segmentation_forests.visualization import plot_segment_comparison

# Compare segment distribution vs background
fig = plot_segment_comparison(
    data=data,
    segment_conditions=[('country', '==', 'UK'), ('device', '==', 'Mobile')],
    metric_column='impressions',
    title='UK Mobile Users vs Background'
)
fig.savefig('segment_comparison.png', dpi=150)

Example output:

The plot shows:

Left: Overlapping histograms (background in blue, segment in coral)
Right: Box plots comparing distributions
Clear separation: Strong segments show minimal overlap

🧠 How It Works

Algorithm Overview

Compute Background Distribution: Calculate the distribution of your metric across all data
Greedy Tree Building:
- At each node, evaluate all feature-value splits
- Choose the split that maximizes divergence from background
- Recursively build left (matching condition) and right (not matching) subtrees
Collect High-Divergence Leaves: Return segments that diverge significantly
Ensemble Aggregation (Forest only): Vote across trees to find robust patterns

Divergence Measures

The algorithm automatically selects the appropriate measure:

Metric Type	Measure	Range	Description
Continuous	Kolmogorov-Smirnov	[0, 1]	Max distance between CDFs
Discrete	Jensen-Shannon	[0, 1]	Symmetric KL divergence

Decision threshold: ≤20 unique values → discrete, >20 → continuous

Quality Guidelines

Interpret divergence scores:

≥ 0.5: 🎯 Excellent - Strong, highly actionable pattern
0.3-0.5: ✓ Good - Meaningful difference worth investigating
0.1-0.3: ⚠️ Weak - Marginal effect, could be noise
< 0.1: ❌ Very weak - Likely statistical noise

📖 API Reference

`SegmentationTree`

SegmentationTree(
    max_depth: int = 5,
    min_samples_split: int = 50,
    min_samples_leaf: int = 20,
    divergence_threshold: float = 0.01,
    random_features: Optional[int] = None
)

Parameters:

max_depth: Maximum tree depth (controls segment complexity)
min_samples_split: Minimum samples required to split a node
min_samples_leaf: Minimum samples required in each child
divergence_threshold: Minimum divergence to keep a segment
random_features: Number of random features per split (None = use all)

Methods:

fit(data: pd.DataFrame, metric_column: str) -> Self: Fit tree to data
get_segments(min_divergence: float = 0.0) -> List[SegmentationNode]: Get segments