fasttps 0.2.0__tar.gz

fasttps-0.2.0/Cargo.lock

@@ -0,0 +1,287 @@
+# This file is automatically @generated by Cargo.
+# It is not intended for manual editing.
+version = 3
+
+[[package]]
+name = "autocfg"
+version = "1.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8"
+
+[[package]]
+name = "cfg-if"
+version = "1.0.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801"
+
+[[package]]
+name = "fasttps"
+version = "0.2.0"
+dependencies = [
+ "fixedbitset",
+ "ndarray",
+ "numpy",
+ "ordered-float",
+ "pyo3",
+]
+
+[[package]]
+name = "fixedbitset"
+version = "0.4.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0ce7134b9999ecaf8bcd65542e436736ef32ddca1b3e06094cb6ec5755203b80"
+
+[[package]]
+name = "heck"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea"
+
+[[package]]
+name = "indoc"
+version = "2.0.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "79cf5c93f93228cf8efb3ba362535fb11199ac548a09ce117c9b1adc3030d706"
+dependencies = [
+ "rustversion",
+]
+
+[[package]]
+name = "libc"
+version = "0.2.182"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6800badb6cb2082ffd7b6a67e6125bb39f18782f793520caee8cb8846be06112"
+
+[[package]]
+name = "matrixmultiply"
+version = "0.3.10"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a06de3016e9fae57a36fd14dba131fccf49f74b40b7fbdb472f96e361ec71a08"
+dependencies = [
+ "autocfg",
+ "rawpointer",
+]
+
+[[package]]
+name = "memoffset"
+version = "0.9.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a"
+dependencies = [
+ "autocfg",
+]
+
+[[package]]
+name = "ndarray"
+version = "0.16.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "882ed72dce9365842bf196bdeedf5055305f11fc8c03dee7bb0194a6cad34841"
+dependencies = [
+ "matrixmultiply",
+ "num-complex",
+ "num-integer",
+ "num-traits",
+ "portable-atomic",
+ "portable-atomic-util",
+ "rawpointer",
+]
+
+[[package]]
+name = "num-complex"
+version = "0.4.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495"
+dependencies = [
+ "num-traits",
+]
+
+[[package]]
+name = "num-integer"
+version = "0.1.46"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f"
+dependencies = [
+ "num-traits",
+]
+
+[[package]]
+name = "num-traits"
+version = "0.2.19"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841"
+dependencies = [
+ "autocfg",
+]
+
+[[package]]
+name = "numpy"
+version = "0.22.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "edb929bc0da91a4d85ed6c0a84deaa53d411abfb387fc271124f91bf6b89f14e"
+dependencies = [
+ "libc",
+ "ndarray",
+ "num-complex",
+ "num-integer",
+ "num-traits",
+ "pyo3",
+ "rustc-hash",
+]
+
+[[package]]
+name = "once_cell"
+version = "1.21.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d"
+
+[[package]]
+name = "ordered-float"
+version = "3.9.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f1e1c390732d15f1d48471625cd92d154e66db2c56645e29a9cd26f4699f72dc"
+dependencies = [
+ "num-traits",
+]
+
+[[package]]
+name = "portable-atomic"
+version = "1.13.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49"
+
+[[package]]
+name = "portable-atomic-util"
+version = "0.2.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7a9db96d7fa8782dd8c15ce32ffe8680bbd1e978a43bf51a34d39483540495f5"
+dependencies = [
+ "portable-atomic",
+]
+
+[[package]]
+name = "proc-macro2"
+version = "1.0.106"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934"
+dependencies = [
+ "unicode-ident",
+]
+
+[[package]]
+name = "pyo3"
+version = "0.22.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f402062616ab18202ae8319da13fa4279883a2b8a9d9f83f20dbade813ce1884"
+dependencies = [
+ "cfg-if",
+ "indoc",
+ "libc",
+ "memoffset",
+ "once_cell",
+ "portable-atomic",
+ "pyo3-build-config",
+ "pyo3-ffi",
+ "pyo3-macros",
+ "unindent",
+]
+
+[[package]]
+name = "pyo3-build-config"
+version = "0.22.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b14b5775b5ff446dd1056212d778012cbe8a0fbffd368029fd9e25b514479c38"
+dependencies = [
+ "once_cell",
+ "target-lexicon",
+]
+
+[[package]]
+name = "pyo3-ffi"
+version = "0.22.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9ab5bcf04a2cdcbb50c7d6105de943f543f9ed92af55818fd17b660390fc8636"
+dependencies = [
+ "libc",
+ "pyo3-build-config",
+]
+
+[[package]]
+name = "pyo3-macros"
+version = "0.22.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0fd24d897903a9e6d80b968368a34e1525aeb719d568dba8b3d4bfa5dc67d453"
+dependencies = [
+ "proc-macro2",
+ "pyo3-macros-backend",
+ "quote",
+ "syn",
+]
+
+[[package]]
+name = "pyo3-macros-backend"
+version = "0.22.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "36c011a03ba1e50152b4b394b479826cad97e7a21eb52df179cd91ac411cbfbe"
+dependencies = [
+ "heck",
+ "proc-macro2",
+ "pyo3-build-config",
+ "quote",
+ "syn",
+]
+
+[[package]]
+name = "quote"
+version = "1.0.44"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "21b2ebcf727b7760c461f091f9f0f539b77b8e87f2fd88131e7f1b433b3cece4"
+dependencies = [
+ "proc-macro2",
+]
+
+[[package]]
+name = "rawpointer"
+version = "0.2.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3"
+
+[[package]]
+name = "rustc-hash"
+version = "1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2"
+
+[[package]]
+name = "rustversion"
+version = "1.0.22"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d"
+
+[[package]]
+name = "syn"
+version = "2.0.117"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "unicode-ident",
+]
+
+[[package]]
+name = "target-lexicon"
+version = "0.12.16"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1"
+
+[[package]]
+name = "unicode-ident"
+version = "1.0.24"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75"
+
+[[package]]
+name = "unindent"
+version = "0.2.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3"

fasttps-0.2.0/Cargo.toml

@@ -0,0 +1,16 @@
+[package]
+name = "fasttps"
+version = "0.2.0"
+edition = "2021"
+readme = "README.md"
+
+[lib]
+name = "fasttps"
+crate-type = ["cdylib"]
+
+[dependencies]
+pyo3 = { version = "0.22", features = ["extension-module"] }
+numpy = "0.22"
+ndarray = "0.16"
+fixedbitset = "0.4"
+ordered-float = "3"

fasttps-0.2.0/PKG-INFO

@@ -0,0 +1,12 @@
+Metadata-Version: 2.4
+Name: fasttps
+Version: 0.2.0
+Summary: Fast TPS internals for the leaguescheduler Python package
+Requires-Python: >=3.13, <3.14
+Description-Content-Type: text/markdown; charset=UTF-8; variant=GFM
+
+# fasttps
+
+Internal Rust extension for [leaguescheduler](https://pypi.org/project/leaguescheduler).
+Not intended for direct use (but be my guest), rather install `leaguescheduler`.
+

fasttps-0.2.0/README.md

@@ -0,0 +1,4 @@
+# fasttps
+
+Internal Rust extension for [leaguescheduler](https://pypi.org/project/leaguescheduler).
+Not intended for direct use (but be my guest), rather install `leaguescheduler`.

fasttps-0.2.0/pyproject.toml

@@ -0,0 +1,13 @@
+[project]
+name = "fasttps"
+version = "0.2.0"
+description = "Fast TPS internals for the leaguescheduler Python package"
+readme = "README.md"
+requires-python = ">=3.13,<3.14"
+
+[build-system]
+requires = ["maturin>=1.0,<2.0"]
+build-backend = "maturin"
+
+[tool.maturin]
+features = ["pyo3/extension-module"]

fasttps-0.2.0/src/lib.rs

@@ -0,0 +1,503 @@
+use std::collections::{HashMap, HashSet};
+
+use fixedbitset::FixedBitSet;
+use ndarray::{Array2, ArrayView2};
+use numpy::{PyArray2, PyArrayMethods, PyReadonlyArray2};
+use ordered_float::OrderedFloat;
+use pyo3::prelude::*;
+use pyo3::types::PyDict;
+
+// ─── Constants ───────────────────────────────────────────────────────────────
+const DISALLOWED_NBR: f64 = 9_999_999.0; // positive if minimization, negative if maximization
+const DISALLOWED_REPLACE: f64 = 1e15;
+const LARGE_NBR: f64 = 9999.0; // must match Python constants.py
+
+// ─── Hungarian algorithm (Kuhn-Munkres) ──────────────────────────────────────
+// Courtesy to https://github.com/neka-nat/fastmunk for main implementation
+fn kuhn_munkres(weights: ArrayView2<f64>, maximize: bool) -> Vec<(usize, usize)> {
+    let nx = weights.nrows();
+    let ny = weights.ncols();
+
+    assert!(
+        nx <= ny,
+        "number of rows must not be larger than number of columns"
+    );
+    let copied_weights = if !maximize {
+        weights.map(|x| -x).view().to_owned()
+    } else {
+        weights.to_owned()
+    };
+    let mut xy: Vec<Option<usize>> = vec![None; nx];
+    let mut yx: Vec<Option<usize>> = vec![None; ny];
+    let mut lx: Vec<OrderedFloat<f64>> = (0..nx)
+        .map(|row| {
+            (0..ny)
+                .map(|col| OrderedFloat(copied_weights[(row, col)]))
+                .max()
+                .unwrap()
+        })
+        .collect::<Vec<_>>();
+    let mut ly: Vec<f64> = vec![0.0; ny];
+    let mut s = FixedBitSet::with_capacity(nx);
+    let mut alternating = Vec::with_capacity(ny);
+    let mut slack = vec![0.0; ny];
+    let mut slackx = Vec::with_capacity(ny);
+    for root in 0..nx {
+        alternating.clear();
+        alternating.resize(ny, None);
+        let mut y = {
+            s.clear();
+            s.insert(root);
+            for y in 0..ny {
+                slack[y] = lx[root].0 + ly[y] - copied_weights[(root, y)];
+            }
+            slackx.clear();
+            slackx.resize(ny, root);
+            Some(loop {
+                let mut delta = f64::INFINITY;
+                let mut x = 0;
+                let mut y = 0;
+                for yy in 0..ny {
+                    if alternating[yy].is_none() && slack[yy] < delta {
+                        delta = slack[yy];
+                        x = slackx[yy];
+                        y = yy;
+                    }
+                }
+                if delta > 0.0 {
+                    for x in s.ones() {
+                        lx[x] = lx[x] - delta;
+                    }
+                    for y in 0..ny {
+                        if alternating[y].is_some() {
+                            ly[y] = ly[y] + delta;
+                        } else {
+                            slack[y] = slack[y] - delta;
+                        }
+                    }
+                }
+                alternating[y] = Some(x);
+                if yx[y].is_none() {
+                    break y;
+                }
+                let x = yx[y].unwrap();
+                s.insert(x);
+                for y in 0..ny {
+                    if alternating[y].is_none() {
+                        let alternate_slack = lx[x] + ly[y] - copied_weights[(x, y)];
+                        if slack[y] > alternate_slack.0 {
+                            slack[y] = alternate_slack.0;
+                            slackx[y] = x;
+                        }
+                    }
+                }
+            })
+        };
+        while y.is_some() {
+            let x = alternating[y.unwrap()].unwrap();
+            let prec = xy[x];
+            yx[y.unwrap()] = Some(x);
+            xy[x] = y;
+            y = prec;
+        }
+    }
+    xy.into_iter()
+        .enumerate()
+        .map(|(i, v)| (i, v.unwrap()))
+        .collect::<Vec<_>>()
+}
+
+// ─── Internal helpers ────────────────────────────────────────────────────────
+
+/// Equivalent of Python legacy _get_team_array(): concatenate X[idx,:] and X[:,idx], filter out LARGE_NBR.
+#[inline]
+fn get_team_array(x: &ArrayView2<f64>, idx: usize) -> Vec<f64> {
+    let n = x.nrows();
+    let mut arr = Vec::with_capacity(2 * n);
+    for j in 0..x.ncols() {
+        let v = x[(idx, j)];
+        if v != LARGE_NBR {
+            arr.push(v);
+        }
+    }
+    for i in 0..n {
+        let v = x[(i, idx)];
+        if v != LARGE_NBR {
+            arr.push(v);
+        }
+    }
+    arr
+}
+
+/// Minimum positive value in a slice of (value - h) differences.
+/// Returns None if no positive value exists.
+#[inline]
+fn min_positive_delta(games: &[f64], h: f64, forward: bool) -> Option<f64> {
+    let mut best = f64::INFINITY;
+    for &g in games {
+        let d = if forward { g - h } else { h - g };
+        if d > 0.0 && d < best {
+            best = d;
+        }
+    }
+    if best == f64::INFINITY {
+        None
+    } else {
+        Some(best)
+    }
+}
+
+/// Build cost matrix for a given team. This is the #1 bottleneck.
+fn create_cost_matrix_inner(
+    x: &ArrayView2<f64>,
+    team_idx: usize,
+    set_home: &[f64],
+    opponents: &[usize],
+    sets_forbidden: &HashMap<usize, HashSet<i64>>,
+    m: f64,
+    r_max: f64,
+    penalties: &[i64],
+    max_penalty_key: usize,
+) -> Array2<f64> {
+    let n_home = set_home.len();
+    let n_oppo = opponents.len();
+    let mut am_cost = Array2::<f64>::zeros((n_home, n_oppo));
+
+    // precompute games for the home team
+    let games_team = get_team_array(x, team_idx);
+
+    // C2 - home date availability (home_dates = set_home)
+
+    // C4 (team part) - precompute which home dates the team already plays on
+    // C5 (team part) - precompute whether any game is within r_max of each home date
+    let mut team_plays_mask = vec![false; n_home];
+    let mut games_in_r_max_team = vec![false; n_home];
+
+    for (i, &h) in set_home.iter().enumerate() {
+        for &g in &games_team {
+            // C4 - team already plays game
+            if g == h {
+                team_plays_mask[i] = true;
+            }
+            // C5 - distance < r_max means too close
+            let d = (g - h).abs() + 1.0;
+            if d < r_max {
+                games_in_r_max_team[i] = true;
+            }
+        }
+    }
+
+    for (j, &oppo_idx) in opponents.iter().enumerate() {
+        let games_oppo = get_team_array(x, oppo_idx);
+
+        // get forbidden set for this opponent (empty if not present)
+        let forbidden_set = sets_forbidden.get(&oppo_idx);
+
+        // C6 - reciprocal game value (game j->i)
+        let reciprocal_val = x[(oppo_idx, team_idx)];
+
+        for (i, &h) in set_home.iter().enumerate() {
+            let h_i64 = h as i64;
+
+            // C3 - forbidden game set
+            let forbidden = match forbidden_set {
+                Some(fs) => fs.contains(&h_i64),
+                None => false,
+            };
+
+            // C4 - team already plays game (precomputed)
+            let team_plays = team_plays_mask[i];
+
+            // C4 - opponent already plays game
+            let mut oppo_plays = false;
+            for &g in &games_oppo {
+                if g == h {
+                    oppo_plays = true;
+                    break;
+                }
+            }
+
+            // C5 - max. 2 games for 'r_max' slots
+            let mut games_in_r_max_oppo = false;
+            for &g in &games_oppo {
+                let d = (g - h).abs() + 1.0;
+                if d < r_max {
+                    games_in_r_max_oppo = true;
+                    break;
+                }
+            }
+            let r_max_violated = games_in_r_max_team[i] || games_in_r_max_oppo;
+
+            // C6 - game i-j is within m days of game j-i
+            let reciprocal_too_close = (h - reciprocal_val).abs() < m;
+
+            // set disallowed cost for all disallowed slots
+            let disallowed =
+                forbidden || team_plays || oppo_plays || reciprocal_too_close || r_max_violated;
+
+            if disallowed {
+                am_cost[(i, j)] = DISALLOWED_NBR;
+            } else {
+                // process penalty calculations for allowed slots
+                let mut total_penalty: i64 = 0;
+
+                // forward-looking for team
+                if let Some(d) = min_positive_delta(&games_team, h, true) {
+                    let key = d as usize;
+                    if key <= max_penalty_key {
+                        total_penalty += penalties[key];
+                    }
+                }
+                // backward-looking for team
+                if let Some(d) = min_positive_delta(&games_team, h, false) {
+                    let key = d as usize;
+                    if key <= max_penalty_key {
+                        total_penalty += penalties[key];
+                    }
+                }
+                // forward-looking for opponent
+                if let Some(d) = min_positive_delta(&games_oppo, h, true) {
+                    let key = d as usize;
+                    if key <= max_penalty_key {
+                        total_penalty += penalties[key];
+                    }
+                }
+                // backward-looking for opponent
+                if let Some(d) = min_positive_delta(&games_oppo, h, false) {
+                    let key = d as usize;
+                    if key <= max_penalty_key {
+                        total_penalty += penalties[key];
+                    }
+                }
+
+                am_cost[(i, j)] = total_penalty as f64;
+            }
+        }
+    }
+
+    am_cost
+}
+
+/// Full solve: cost matrix → adjacency matrix → Munkres → process indexes → (pick, total_cost).
+fn solve_inner(
+    x: &ArrayView2<f64>,
+    team_idx: usize,
+    set_home: &[f64],
+    opponents: &[usize],
+    sets_forbidden: &HashMap<usize, HashSet<i64>>,
+    m: f64,
+    p: f64,
+    r_max: f64,
+    penalties: &[i64],
+    max_penalty_key: usize,
+) -> (Vec<f64>, f64) {
+    let n_home = set_home.len();
+    let n_oppo = opponents.len();
+    let dim = n_home + n_oppo;
+
+    // construct cost matrix
+    let am_cost = create_cost_matrix_inner(
+        x,
+        team_idx,
+        set_home,
+        opponents,
+        sets_forbidden,
+        m,
+        r_max,
+        penalties,
+        max_penalty_key,
+    );
+
+    // construct full adjacency matrix:
+    // [ am_cost  |  zeros ]
+    // [ am_bott  |  zeros ]
+    // where am_bott is n_oppo x n_oppo filled with p
+    // and zeros is dim x (dim - n_oppo)
+    let mut am = Array2::<f64>::zeros((dim, dim));
+
+    // fill cost block (top-left: n_home x n_oppo)
+    for i in 0..n_home {
+        for j in 0..n_oppo {
+            let v = am_cost[(i, j)];
+            // replace disallowed values with a large number
+            am[(i, j)] = if v == DISALLOWED_NBR {
+                DISALLOWED_REPLACE
+            } else {
+                v
+            };
+        }
+    }
+
+    // fill bottom block (n_oppo x n_oppo) with p
+    for i in 0..n_oppo {
+        for j in 0..n_oppo {
+            am[(n_home + i, j)] = p;
+        }
+    }
+    // right block stays zeros (already initialized)
+
+    // run Hungarian algorithm (minimize)
+    let indexes = kuhn_munkres(am.view(), false);
+
+    // compute total cost
+    let total_cost: f64 = indexes.iter().map(|&(r, c)| am[(r, c)]).sum();
+
+    // process optimal indexes: sort by opponent, map to home slot or NaN
+    let mut indexes_inv: Vec<(usize, f64)> = Vec::with_capacity(n_oppo);
+    for &(k, v) in &indexes {
+        if v < n_oppo {
+            let slot = if k < n_home { set_home[k] } else { f64::NAN };
+            indexes_inv.push((opponents[v], slot));
+        }
+    }
+    indexes_inv.sort_by_key(|&(opp, _)| opp);
+
+    let pick: Vec<f64> = indexes_inv.iter().map(|&(_, slot)| slot).collect();
+
+    (pick, total_cost)
+}
+
+// ─── PyO3 class ──────────────────────────────────────────────────────────────
+
+#[pyclass(module = "fasttps")]
+struct FastTPS {
+    sets_home: HashMap<usize, Vec<f64>>,
+    sets_forbidden: HashMap<usize, HashSet<i64>>,
+    m: f64,
+    p: f64,
+    r_max: f64,
+    /// Flat penalty lookup: penalties_vec[d] = penalty for distance d.
+    penalties_vec: Vec<i64>,
+    max_penalty_key: usize,
+}
+
+#[pymethods]
+impl FastTPS {
+    #[new]
+    #[pyo3(signature = (sets_home, sets_forbidden, m, p, r_max, penalties))]
+    fn new(
+        sets_home: &Bound<'_, PyDict>,
+        sets_forbidden: &Bound<'_, PyDict>,
+        m: i64,
+        p: i64,
+        r_max: i64,
+        penalties: &Bound<'_, PyDict>,
+    ) -> PyResult<Self> {
+        // parse sets_home: dict[int, list/set of float/int]
+        let mut sh: HashMap<usize, Vec<f64>> = HashMap::new();
+        for (k, v) in sets_home.iter() {
+            let key: usize = k.extract()?;
+            let vals: Vec<f64> = v.extract()?;
+            sh.insert(key, vals);
+        }
+
+        // parse sets_forbidden: dict[int, list/set of float/int] → HashSet<i64>
+        let mut sf: HashMap<usize, HashSet<i64>> = HashMap::new();
+        for (k, v) in sets_forbidden.iter() {
+            let key: usize = k.extract()?;
+            let vals: Vec<i64> = v.extract()?;
+            sf.insert(key, vals.into_iter().collect());
+        }
+
+        // parse penalties dict to flat Vec
+        let mut max_key: usize = 0;
+        let mut pen_map: HashMap<usize, i64> = HashMap::new();
+        for (k, v) in penalties.iter() {
+            let key: usize = k.extract()?;
+            let val: i64 = v.extract()?;
+            pen_map.insert(key, val);
+            if key > max_key {
+                max_key = key;
+            }
+        }
+        let mut penalties_vec = vec![0i64; max_key + 1];
+        for (k, v) in &pen_map {
+            penalties_vec[*k] = *v;
+        }
+
+        Ok(FastTPS {
+            sets_home: sh,
+            sets_forbidden: sf,
+            m: m as f64,
+            p: p as f64,
+            r_max: r_max.max(2) as f64,
+            penalties_vec,
+            max_penalty_key: max_key,
+        })
+    }
+
+    /// Solve transportation problem for given home team.
+    /// Modifies X in-place and returns total_cost.
+    fn solve<'py>(
+        &self,
+        py: Python<'py>,
+        x_py: &Bound<'py, PyArray2<f64>>,
+        team_idx: usize,
+    ) -> PyResult<f64> {
+        let x_read = x_py.readonly();
+        let x_view = x_read.as_array();
+        let n = x_view.nrows();
+
+        let set_home = self
+            .sets_home
+            .get(&team_idx)
+            .expect("team_idx not in sets_home");
+        let opponents: Vec<usize> = (0..n).filter(|&t| t != team_idx).collect();
+
+        let (pick, total_cost) = solve_inner(
+            &x_view,
+            team_idx,
+            set_home,
+            &opponents,
+            &self.sets_forbidden,
+            self.m,
+            self.p,
+            self.r_max,
+            &self.penalties_vec,
+            self.max_penalty_key,
+        );
+
+        // release readonly borrow before writing
+        drop(x_read);
+
+        // assign selection to X in-place: X[team_idx, opponents] = pick
+        unsafe {
+            let mut x_mut = x_py.as_array_mut();
+            for (idx, &opp) in opponents.iter().enumerate() {
+                *x_mut.uget_mut([team_idx, opp]) = pick[idx];
+            }
+        }
+
+        Ok(total_cost)
+    }
+
+    /// Create cost matrix (exposed for construction phase method 2).
+    fn create_cost_matrix<'py>(
+        &self,
+        py: Python<'py>,
+        x_py: PyReadonlyArray2<'py, f64>,
+        team_idx: usize,
+        set_home: Vec<f64>,
+        opponents: Vec<usize>,
+    ) -> PyResult<Py<PyArray2<f64>>> {
+        let x = x_py.as_array();
+        let am = create_cost_matrix_inner(
+            &x,
+            team_idx,
+            &set_home,
+            &opponents,
+            &self.sets_forbidden,
+            self.m,
+            self.r_max,
+            &self.penalties_vec,
+            self.max_penalty_key,
+        );
+        Ok(PyArray2::from_owned_array_bound(py, am).unbind())
+    }
+}
+
+#[pymodule]
+fn fasttps(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    m.add_class::<FastTPS>()?;
+    Ok(())
+}