seq-len-balance 0.1.0__tar.gz

seq_len_balance-0.1.0/.gitignore

@@ -0,0 +1,30 @@
+# Generated by Cargo
+# will have compiled files and executables
+debug
+target
+
+# These are backup files generated by rustfmt
+**/*.rs.bk
+
+# MSVC Windows builds of rustc generate these, which store debugging information
+*.pdb
+
+# Generated by cargo mutants
+# Contains mutation testing data
+**/mutants.out*/
+
+# RustRover
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+
+# Added by cargo
+
+/target
+
+.ipynb_checkpoints/
+
+__pycache__

seq_len_balance-0.1.0/.python-version

@@ -0,0 +1 @@
+3.11.3

seq_len_balance-0.1.0/Cargo.lock

@@ -0,0 +1,158 @@
+# This file is automatically @generated by Cargo.
+# It is not intended for manual editing.
+version = 4
+
+[[package]]
+name = "autocfg"
+version = "1.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8"
+
+[[package]]
+name = "heck"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea"
+
+[[package]]
+name = "libc"
+version = "0.2.182"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6800badb6cb2082ffd7b6a67e6125bb39f18782f793520caee8cb8846be06112"
+
+[[package]]
+name = "num-traits"
+version = "0.2.19"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841"
+dependencies = [
+ "autocfg",
+]
+
+[[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 = "5.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7f4779c6901a562440c3786d08192c6fbda7c1c2060edd10006b05ee35d10f2d"
+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 = "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.28.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cf85e27e86080aafd5a22eae58a162e133a589551542b3e5cee4beb27e54f8e1"
+dependencies = [
+ "libc",
+ "once_cell",
+ "portable-atomic",
+ "pyo3-build-config",
+ "pyo3-ffi",
+ "pyo3-macros",
+]
+
+[[package]]
+name = "pyo3-build-config"
+version = "0.28.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8bf94ee265674bf76c09fa430b0e99c26e319c945d96ca0d5a8215f31bf81cf7"
+dependencies = [
+ "target-lexicon",
+]
+
+[[package]]
+name = "pyo3-ffi"
+version = "0.28.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "491aa5fc66d8059dd44a75f4580a2962c1862a1c2945359db36f6c2818b748dc"
+dependencies = [
+ "libc",
+ "pyo3-build-config",
+]
+
+[[package]]
+name = "pyo3-macros"
+version = "0.28.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f5d671734e9d7a43449f8480f8b38115df67bef8d21f76837fa75ee7aaa5e52e"
+dependencies = [
+ "proc-macro2",
+ "pyo3-macros-backend",
+ "quote",
+ "syn",
+]
+
+[[package]]
+name = "pyo3-macros-backend"
+version = "0.28.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "22faaa1ce6c430a1f71658760497291065e6450d7b5dc2bcf254d49f66ee700a"
+dependencies = [
+ "heck",
+ "proc-macro2",
+ "pyo3-build-config",
+ "quote",
+ "syn",
+]
+
+[[package]]
+name = "quote"
+version = "1.0.45"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924"
+dependencies = [
+ "proc-macro2",
+]
+
+[[package]]
+name = "seq-len-balance-rs"
+version = "0.1.0"
+dependencies = [
+ "ordered-float",
+ "pyo3",
+]
+
+[[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.13.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca"
+
+[[package]]
+name = "unicode-ident"
+version = "1.0.24"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75"

seq_len_balance-0.1.0/Cargo.toml

@@ -0,0 +1,17 @@
+[package]
+name = "seq-len-balance-rs"
+version = "0.1.0"
+edition = "2021"
+readme = "README.md"
+
+[lib]
+name = "seq_len_balance"
+crate-type = ["cdylib", "rlib"]
+
+[[bin]]
+name = "seq-len-balance-rs"
+path = "src/main.rs"
+
+[dependencies]
+pyo3 = { version = "0.28", features = ["extension-module"] }
+ordered-float = "5"

seq_len_balance-0.1.0/Makefile

@@ -0,0 +1,40 @@
+# ---- Project metadata ----
+PACKAGE_NAME := seq-len-balance-rs
+PYTHON       ?= python3
+BUILD_DIR    := target/wheels
+
+# ---- Commands ----
+
+.PHONY: help build develop publish clean
+
+help:
+	@echo "Available commands:"
+	@echo "  make build     - Build release wheels (and sdist)"
+	@echo "  make develop   - Install locally in dev mode"
+	@echo "  make publish   - Publish to PyPI using maturin"
+	@echo "  make clean     - Remove build artifacts"
+
+build:
+	@echo "🚀 Building release wheels..."
+	maturin build --release --sdist
+	@echo "✅ Wheels generated in $(BUILD_DIR)/"
+
+develop:
+	@echo "💡 Installing in development mode..."
+	maturin develop
+
+publish: build
+	@echo "📦 Publishing $(PACKAGE_NAME) to PyPI..."
+	@if [ -z "$$MATURIN_PASSWORD" ]; then \
+		echo "❌ Error: MATURIN_PASSWORD (your PyPI token) is not set."; \
+		echo "   export MATURIN_USERNAME=__token__"; \
+		echo "   export MATURIN_PASSWORD=pypi-xxxxxx"; \
+		exit 1; \
+	fi
+	maturin publish --skip-existing
+	@echo "✅ Published successfully!"
+
+clean:
+	@echo "🧹 Cleaning build artifacts..."
+	cargo clean
+	rm -rf $(BUILD_DIR) dist build *.egg-info

seq_len_balance-0.1.0/PKG-INFO

@@ -0,0 +1,5 @@
+Metadata-Version: 2.4
+Name: seq-len-balance
+Version: 0.1.0
+Summary: Fast sequence-length balancing: KK partition, FFD, and BFD bin packing
+Requires-Python: >=3.11

seq_len_balance-0.1.0/README.md

@@ -0,0 +1,48 @@
+# seq-len-balance-rs
+
+Fast sequence-length balancing: KK partition, FFD, and BFD bin packing (Rust + Python).
+
+## Setup
+
+```bash
+uv sync
+pip install maturin
+maturin develop
+```
+
+## Usage
+
+```python
+from seq_len_balance import kk, ffd, bfd
+
+print(kk([8, 7, 6, 5, 4], 2))   # -> [[8, 5], [7, 4], [6]]
+print(ffd([6, 3, 4, 5, 2, 7, 1], 10))  # -> [[7, 3], [6, 4], [5, 2, 1]]
+print(bfd([6, 3, 4, 5, 2, 7, 1], 10))  # -> [[7, 2, 1], [6, 4], [5, 3]]
+```
+
+## Benchmark
+
+```bash
+uv sync --group benchmark
+uv run python test_kk.py
+```
+
+```
+=================================================================
+KK (greedy LPT) benchmark  k=4  repeats=3
+=================================================================
+n=   10,000  Python:      2.4 ms   Rust:      0.4 ms
+n=  100,000  Python:     25.5 ms   Rust:      2.9 ms
+n=  500,000  Python:    131.4 ms   Rust:     16.3 ms
+n=1,000,000  Python:    294.7 ms   Rust:     33.1 ms
+n=2,000,000  Python:    663.9 ms   Rust:     69.6 ms
+n=4,000,000  Python:   1405.9 ms   Rust:    144.9 ms
+n=10,000,000  Python:   3577.8 ms   Rust:    377.6 ms
+```
+
+## Rust in Jupyter (optional)
+
+```bash
+cargo install evcxr_jupyter
+evcxr_jupyter --install
+```

seq_len_balance-0.1.0/pyproject.toml

@@ -0,0 +1,26 @@
+[build-system]
+requires = ["maturin>=1.0,<2.0"]
+build-backend = "maturin"
+
+[project]
+name = "seq-len-balance"
+version = "0.1.0"
+requires-python = ">=3.11"
+description = "Fast sequence-length balancing: KK partition, FFD, and BFD bin packing"
+dependencies = []
+
+[tool.maturin]
+features = ["pyo3/extension-module"]
+
+[dependency-groups]
+benchmark = [
+    "matplotlib>=3.10.8",
+    "pytest>=9.0.2",
+]
+
+[tool.uv]
+dev-dependencies = [
+    "matplotlib>=3.10.8",
+    "maturin>=1.0,<2.0",
+    "numpy>=2.4.2",
+]

seq_len_balance-0.1.0/src/bin.rs

@@ -0,0 +1,173 @@
+use std::fmt;
+use std::ops::{Add, Sub};
+
+/// Unified result of a bin packing or partitioning algorithm.
+///
+/// - `capacity`: `Some(limit)` for bin packing (FFD/BFD), `None` for unconstrained partitioning (KK).
+/// - `remaining[i]`: unused space in bin i; zero for partitioning results.
+/// - `sums[i]`: total item weight assigned to bin/group i.
+#[derive(Debug)]
+pub struct BinPacking<T> {
+    pub bins: Vec<Vec<T>>,
+    pub sums: Vec<T>,
+    pub remaining: Vec<T>,
+    pub capacity: Option<T>,
+}
+
+impl<T> BinPacking<T>
+where
+    T: Ord + Clone + Sub<Output = T>,
+{
+    /// Number of bins / groups used.
+    pub fn num_bins(&self) -> usize {
+        self.bins.len()
+    }
+
+    /// Difference between the largest and smallest bin sum.
+    pub fn imbalance(&self) -> T {
+        let max = self.sums.iter().max().unwrap().clone();
+        let min = self.sums.iter().min().unwrap().clone();
+        max - min
+    }
+
+    /// Total wasted space across all bins (sum of `remaining`).
+    /// For partitioning results this is always zero.
+    pub fn waste(&self) -> T
+    where
+        T: Default + Add<Output = T>,
+    {
+        self.remaining
+            .iter()
+            .cloned()
+            .fold(T::default(), |acc, r| acc + r)
+    }
+}
+
+impl<T: fmt::Display> fmt::Display for BinPacking<T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        for (i, ((bin, sum), rem)) in self
+            .bins
+            .iter()
+            .zip(self.sums.iter())
+            .zip(self.remaining.iter())
+            .enumerate()
+        {
+            match &self.capacity {
+                Some(_) => write!(f, "Bin {i} (sum={sum}, remaining={rem}): [")?,
+                None => write!(f, "Bin {i} (sum={sum}): [")?,
+            }
+            for (j, item) in bin.iter().enumerate() {
+                if j > 0 {
+                    write!(f, ", ")?;
+                }
+                write!(f, "{item}")?;
+            }
+            writeln!(f, "]")?;
+        }
+        Ok(())
+    }
+}
+
+/// First Fit Decreasing bin packing.
+///
+/// Sorts items descending, then for each item tries bins in order (0, 1, 2, ...)
+/// and places it in the **first** bin with enough remaining capacity.
+/// Opens a new bin only when no existing bin fits.
+/// Capacity defaults to the sum of all items when not provided.
+pub fn first_fit_decreasing<T>(mut items: Vec<T>, capacity: Option<T>) -> BinPacking<T>
+where
+    T: Ord + Add<Output = T> + Sub<Output = T> + Default + Clone,
+{
+    let capacity = capacity
+        .unwrap_or_else(|| items.iter().cloned().fold(T::default(), |acc, x| acc + x));
+
+    items.sort_unstable_by(|a, b| b.cmp(a));
+
+    let mut bins: Vec<Vec<T>> = Vec::new();
+    let mut sums: Vec<T> = Vec::new();
+    let mut remaining: Vec<T> = Vec::new();
+
+    for item in items {
+        assert!(item <= capacity, "item exceeds bin capacity");
+
+        // First bin with enough room.
+        let target = bins
+            .iter()
+            .enumerate()
+            .find(|(i, _)| remaining[*i] >= item)
+            .map(|(i, _)| i);
+
+        match target {
+            Some(i) => {
+                remaining[i] = remaining[i].clone() - item.clone();
+                sums[i] = sums[i].clone() + item.clone();
+                bins[i].push(item);
+            }
+            None => {
+                remaining.push(capacity.clone() - item.clone());
+                sums.push(item.clone());
+                bins.push(vec![item]);
+            }
+        }
+    }
+
+    BinPacking {
+        bins,
+        sums,
+        remaining,
+        capacity: Some(capacity),
+    }
+}
+
+/// Best Fit Decreasing bin packing.
+///
+/// Sorts items descending, then for each item finds the bin with the
+/// **minimum remaining capacity** that still fits the item (tightest fit).
+/// This leaves larger gaps available for larger future items.
+/// Opens a new bin only when no existing bin fits.
+/// Capacity defaults to the sum of all items when not provided.
+pub fn best_fit_decreasing<T>(mut items: Vec<T>, capacity: Option<T>) -> BinPacking<T>
+where
+    T: Ord + Add<Output = T> + Sub<Output = T> + Default + Clone,
+{
+    let capacity = capacity
+        .unwrap_or_else(|| items.iter().cloned().fold(T::default(), |acc, x| acc + x));
+
+    items.sort_unstable_by(|a, b| b.cmp(a));
+
+    let mut bins: Vec<Vec<T>> = Vec::new();
+    let mut sums: Vec<T> = Vec::new();
+    let mut remaining: Vec<T> = Vec::new();
+
+    for item in items {
+        assert!(item <= capacity, "item exceeds bin capacity");
+
+        // Bin with the smallest remaining space that still fits (tightest fit).
+        let best_idx = remaining
+            .iter()
+            .enumerate()
+            .filter(|(_, r)| *r >= &item)
+            .min_by(|(_, a), (_, b)| a.cmp(b))
+            .map(|(i, _)| i);
+
+        match best_idx {
+            Some(i) => {
+                remaining[i] = remaining[i].clone() - item.clone();
+                sums[i] = sums[i].clone() + item.clone();
+                bins[i].push(item);
+            }
+            None => {
+                remaining.push(capacity.clone() - item.clone());
+                sums.push(item.clone());
+                bins.push(vec![item]);
+            }
+        }
+    }
+
+    BinPacking {
+        bins,
+        sums,
+        remaining,
+        capacity: Some(capacity),
+    }
+}

seq_len_balance-0.1.0/src/heapq.rs

@@ -0,0 +1,116 @@
+#[derive(Debug, Default)]
+pub struct MinHeap<T: Ord> {
+    data: Vec<T>,
+}
+
+impl<T: Ord> MinHeap<T> {
+    pub fn len(&self) -> usize {
+        self.data.len()
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.data.is_empty()
+    }
+    // used after inserting a new element
+    pub fn sift_up(&mut self, mut i: usize) {
+        while i > 0 {
+            let parent = (i - 1) / 2;
+
+            if self.data[parent] <= self.data[i] {
+                break;
+            }
+            self.data.swap(parent, i);
+            i = parent;
+        }
+    }
+    // after deleting an element
+    pub fn sift_down(&mut self, mut i: usize) {
+        let n = self.data.len();
+
+        loop {
+            // parent(i) = (i - 1) / 2
+            // left(i)   = 2*i + 1
+            // right(i)  = 2*i + 2
+            let left = 2 * i + 1;
+            let right = 2 * i + 2;
+
+            if left >= n {
+                break;
+            }
+            let mut smallest = left;
+            if right < n && self.data[right] < self.data[left] {
+                smallest = right;
+            }
+
+            if self.data[i] <= self.data[smallest] {
+                break;
+            }
+            self.data.swap(i, smallest);
+            i = smallest;
+        }
+    }
+
+    pub fn push(&mut self, val: T) {
+        self.data.push(val);
+        let i = self.data.len() - 1;
+        self.sift_up(i);
+    }
+
+    pub fn pop(&mut self) -> Option<T> {
+        if self.data.is_empty() {
+            return None;
+        }
+        let last = self.data.len() - 1;
+        self.data.swap(0, last);
+        let top = self.data.pop();
+        if !self.data.is_empty() {
+            self.sift_down(0);
+        }
+        top
+    }
+
+    pub fn peek(&self) -> Option<&T> {
+        self.data.first()
+    }
+
+    pub fn from_vec(data: Vec<T>) -> Self {
+        let mut h = Self { data };
+        let n = h.data.len();
+        for i in (0..n / 2).rev() {
+            h.sift_down(i);
+        }
+        h
+    }
+
+    pub fn with_capacity(cap: usize) -> Self {
+        Self {
+            data: Vec::with_capacity(cap),
+        }
+    }
+}
+
+impl<T: Ord + std::fmt::Display> MinHeap<T> {
+    pub fn print_tree(&self) {
+        if self.data.is_empty() {
+            println!("(empty)");
+            return;
+        }
+
+        let n = self.data.len();
+        let mut level_start = 0;
+        let mut level_size = 1;
+
+        while level_start < n {
+            let level_end = (level_start + level_size).min(n);
+            for i in level_start..level_end {
+                if i > level_start {
+                    print!("  ");
+                }
+                print!("{}", self.data[i]);
+            }
+            println!();
+            level_start += level_size;
+            level_size *= 2;
+        }
+    }
+}

seq_len_balance-0.1.0/src/kk.rs

@@ -0,0 +1,56 @@
+use crate::bin::BinPacking;
+use std::cmp::Reverse;
+use std::collections::BinaryHeap;
+use std::ops::Add;
+
+// [8,7,6,5,4]
+// [0 | 0 | 8] -> sort ascending (first element)
+// [7 | 0 | 0] -> sort descending element
+// [7 | 0 | 8] -> sort -> [0 | 7 | 8]
+
+// [0 | 7 | 8] -> sorted ascending
+// [6 | 0 | 0] -> sorted descending
+// [6 | 7 | 8] -> sort ascending -> [6 | 7 | 8]
+
+// [6 | 7 | 8]
+// [5 | 0 | 0]
+// [(6,5) | 7 | 8] -> sort ascending -> [7 | 8 | 11]
+
+// [7 | 8 | 11]
+// [5 | 0 | 0]
+// [(7,5) | 8 | (6,5)] -> [8 | 11 | 12]
+
+// [8 | 11 | 12]
+// [4 | 0 | 0]
+// [(8,4) | 11 | 12] -> [(8,4) | (6,5) | (7,5) ]
+
+pub fn kk_partition<T>(mut items: Vec<T>, k: usize) -> BinPacking<T>
+where
+    T: Ord + Add<Output = T> + Default + Clone,
+{
+    assert!(k > 0, "k must be at least 1");
+    // placing large items first improves balancing.
+    items.sort_unstable_by(|a, b| b.cmp(a));
+
+    let mut bins: Vec<Vec<T>> = vec![vec![]; k];
+    let mut sums: Vec<T> = vec![T::default(); k];
+    // BinaryHeap<Reverse<...>> gives min-heap behaviour.
+    let mut min_heap: BinaryHeap<Reverse<(T, usize)>> =
+        (0..k).map(|i| Reverse((T::default(), i))).collect();
+
+    for item in items {
+        let Reverse((min_sum, idx)) = min_heap.pop().unwrap();
+        let new_sum = min_sum + item.clone();
+        bins[idx].push(item);
+        sums[idx] = new_sum.clone();
+        min_heap.push(Reverse((new_sum, idx)));
+    }
+
+    let remaining = vec![T::default(); k];
+    BinPacking {
+        bins,
+        sums,
+        remaining,
+        capacity: None,
+    }
+}