segmtree

Interpretable unsupervised decision trees for segmentation — clusters rows of a tabular dataset into segments, where every segment is described by a short, human-readable rule chain. No target variable required.

segmtree implements a clustering tree: instead of predicting labels, it recursively splits the data so that within-segment heterogeneity drops as much as possible at every cut, growing best-first (highest-gain node expanded first). Think of it as "a decision tree grown without y" — you get both the partition and the rules that define it, which plain k-means or GMM cannot give you.

Why segmtree?

Typical uses: customer segmentation with behavioral flags, patient phenotyping, survey respondent typing — any "many 0/1 tags + a few counts" table.

Installation

pip install segmtree

Requires Python ≥ 3.9 and numpy only.

Quickstart

import numpy as np
from segmtree import SegTree, extract_rules, profile_leaves, replay_check, bootstrap_stability

rng = np.random.default_rng(0)
# 4 continuous features + 3 binary flags, three latent groups
X = np.vstack([
    np.hstack([rng.normal([0, -1, 2, .5], .5, (250, 4)), (rng.random((250, 3)) < .1)]),
    np.hstack([rng.normal([4, 3, -2, .5], .5, (200, 4)), (rng.random((200, 3)) < .5)]),
    np.hstack([rng.normal([-4, 2, 0, 8], .5, (150, 4)), (rng.random((150, 3)) < .9)]),
])

tree = SegTree(min_gain=0.02, min_leaf_frac=0.05, min_leaf_abs=10)
tree.fit(X)

print(f"{tree.n_leaves_} segments, heterogeneity reduced {tree.reduction_:.0%}")

for rule in extract_rules(tree):
    print(rule)                      # #0: x1 <= -1.24 AND x0 <= 1.83 ...

profile_leaves(tree, X, top_k=3)     # most distinctive features per segment
replay_check(tree, X)                # True: stored rules reproduce the partition exactly
bootstrap_stability(tree, X, n_replicates=12, random_state=0).mean   # e.g. 0.98

Binary columns (all values in {0, 1}) are detected automatically and split with a single test; pandas DataFrames are accepted and column names flow into the rules.

Categorical features

Columns holding integer category codes (e.g. 0..k-1 from a nominal encoding) can be declared with categorical_features. They are split by the best subset of categories rather than by a threshold, and the rules read as membership tests:

tree = SegTree(categorical_features=["color"], max_categories=12).fit(df)
# rule example:  #3: color in {1, 2} AND income > 42 ...

Impurity contribution: categorical columns use Gini impurity relative to the root Gini, while continuous/binary columns keep normalized variance. Gini depends only on the category distribution, so any one-to-one recoding of the codes yields an equivalent partition — results are encoding-invariant.

Notes:

• Codes must be integral; non-numeric labels should be factorized first (pd.factorize, df["color"].astype("category").cat.codes).
• Codes are handled exactly at any magnitude (int64/float64 internally), even with dtype=np.float32 for the continuous columns.
• All 2^(k-1)-1 category partitions of a node are evaluated exactly, so keep k <= max_categories (rare levels can be grouped beforehand); for high-cardinality columns prefer one-hot encoding.
• Unseen codes at prediction time follow handle_unknown: "complement" (default) routes them deterministically to the complement side and replay_check stays exact; "error" raises ValueError naming the feature and unknown values. The training-time code sets are exposed via the categories_ attribute.
• profile_leaves reports categorical structure through category_enrichment — (feature, code, leaf_share, lift) pairs — and keeps z-scores for continuous/binary features only.

How it works

Node heterogeneity averages per-feature variance relative to the root:

H(S) = (1/m) · Σ_j Var_j(S) / Var_j(root)          H(root) = 1

Because Bernoulli variance p(1−p) fits the same formula, binary and continuous features share one comparable criterion. A split into L/R earns:

Gain = H(S) − [ n_L/n_S · H(L) + n_R/n_S · H(R) ]

Candidates are the in-node quantiles of each feature (quantiles=8 → the 5% … 95% quantiles; binary features get one candidate at 0.5). All candidates are evaluated per feature via sort + prefix sums in O(n log n), and the globally best leaf is always split next (best-first).

Splitting stops when no cut gains at least min_gain, a node is too small (max(min_leaf_abs, min_leaf_frac·n) per child), or max_depth is reached.

Two properties make results trustworthy out of the box:

• Gain additivity — total_gain_ == h_root_ − h_avg_ up to float error; it accounts for exactly how much heterogeneity was explained.
• Rule replay — replay_check() re-evaluates every leaf's rule chain from scratch and must reproduce predict() bit-for-bit.

The test suite additionally verifies exact agreement against an independent naive reference implementation.

API overview

SegTree(min_gain=0.01, min_leaf_frac=0.01, min_leaf_abs=50, max_depth=None,
        quantiles=8, binary_features="auto", categorical_features=None,
        max_categories=12, handle_unknown="complement", feature_names=None,
        dtype=np.float64)

Tuning tips:

• Fewer, larger segments → raise min_gain (e.g. 0.02–0.05) or lower quantiles; cap max_depth (e.g. 4–6) to keep rules short and stable.
• Small datasets (< 5000 rows) → lower min_leaf_abs (it defaults to 50).
• Unstable segments (low bootstrap ARI) → fewer quantiles, higher min_gain, or a depth cap usually stabilize the structure.

Comparison with related tools

• scikit-learn DecisionTree on k-means labels ("surrogate tree"): two-step, rules may contradict the clustering; segmtree optimizes split quality and rule fidelity simultaneously.
• forest-clustering / URF: forest-proximity clustering — accurate but the segments themselves are not directly interpretable.
• CUBT (R): closest academic relative; segmtree adds best-first growth, mixed-type normalization, additivity checks and bootstrap stability in a numpy-only package.

Development

git clone <your-fork-url> && cd segmtree
pip install -e .[dev]
pytest            # run the test suite
ruff check src tests

License

MIT — see LICENSE.

中文文档