From 19ee178d183b858fec9ae4a2782c6ccffa76cbfc Mon Sep 17 00:00:00 2001
From: Kevin Thornton <krthornt@uci.edu>
Date: Tue, 5 May 2026 12:39:43 -0700
Subject: [PATCH] perf: specialize nth for tree sequence individual iterator

---
 src/sys/treeseq.rs   | 6 ++++++
 src/trees/treeseq.rs | 5 +++++
 tests/test_trees.rs  | 3 +++
 3 files changed, 14 insertions(+)

diff --git a/src/sys/treeseq.rs b/src/sys/treeseq.rs
index 2c9488f3..513219c8 100644
--- a/src/sys/treeseq.rs
+++ b/src/sys/treeseq.rs
@@ -31,6 +31,12 @@ impl<'ts> Iterator for TreeSeqIndividualIter<'ts> {
         self.current_row += 1;
         self.treeseq.individual(r)
     }
+
+    fn nth(&mut self, n: usize) -> Option<Self::Item> {
+        let u = bindings::tsk_id_t::try_from(n).ok()?;
+        self.current_row = u + 1;
+        self.treeseq.individual(u)
+    }
 }
 
 impl TreeSequence {
diff --git a/src/trees/treeseq.rs b/src/trees/treeseq.rs
index 22cef9cf..3abbb95d 100644
--- a/src/trees/treeseq.rs
+++ b/src/trees/treeseq.rs
@@ -547,6 +547,11 @@ impl TreeSequence {
     }
 
     /// Return an iterator over the individuals.
+    ///
+    /// # Notes
+    ///
+    /// * [`Iterator::nth`] is specialized for the return value,
+    ///   making `nth` and [`Iterator::skip`] constant-time.
     pub fn individual_iter(&self) -> impl Iterator<Item = crate::Individual<'_>> {
         self.inner.individual_iter()
     }
diff --git a/tests/test_trees.rs b/tests/test_trees.rs
index 304e42a7..fa9fd10e 100644
--- a/tests/test_trees.rs
+++ b/tests/test_trees.rs
@@ -919,6 +919,9 @@ fn test_treeseq_individual_iter() {
             assert_eq!(location, [1., 2.]);
         }
     }
+    for (i, ind) in ts.individual_iter().enumerate() {
+        assert_eq!(ts.individual_iter().nth(i).unwrap(), ind)
+    }
 }
 
 #[test]