src/color_group.rs

@@ -46,6 +46,19 @@ impl ColorGroups {
     pub(crate) fn total_events(&self) -> usize {
         self.groups.iter().map(|g| g.len()).sum()
     }
+
+    /// Contiguous index range for one color after events have been reordered
+    /// into color-contiguous positions by `TimeSlice::recompute_color_groups`.
+    #[allow(dead_code)]
+    pub(crate) fn color_range(&self, color_idx: usize) -> std::ops::Range<usize> {
+        let group = &self.groups[color_idx];
+        if group.is_empty() {
+            return 0..0;
+        }
+        let start = *group.first().unwrap();
+        let end = *group.last().unwrap() + 1;
+        start..end
+    }
 }
 
 /// Compute color groups greedily.

src/time_slice.rs

@@ -7,6 +7,7 @@ use std::collections::HashMap;
 use crate::{
     Index, N_INF,
     arena::ScratchArena,
+    color_group::ColorGroups,
     drift::Drift,
     game::Game,
     gaussian::Gaussian,
@@ -84,6 +85,12 @@ pub(crate) struct Event {
 }
 
 impl Event {
+    pub(crate) fn iter_agents(&self) -> impl Iterator<Item = Index> + '_ {
+        self.teams
+            .iter()
+            .flat_map(|t| t.items.iter().map(|it| it.agent))
+    }
+
     fn outputs(&self) -> Vec<f64> {
         self.teams
             .iter()
@@ -117,6 +124,7 @@ pub struct TimeSlice<T: Time = i64> {
     pub(crate) time: T,
     p_draw: f64,
     arena: ScratchArena,
+    pub(crate) color_groups: ColorGroups,
 }
 
 impl<T: Time> TimeSlice<T> {
@@ -127,9 +135,44 @@ impl<T: Time> TimeSlice<T> {
             time,
             p_draw,
             arena: ScratchArena::new(),
+            color_groups: ColorGroups::new(),
         }
     }
 
+    /// Recompute the color-group partition and reorder `self.events` into
+    /// color-contiguous ranges. After this call, `self.color_groups.groups[c]`
+    /// contains a contiguous ascending range of indices in `self.events`.
+    pub(crate) fn recompute_color_groups(&mut self) {
+        use crate::color_group::color_greedy;
+
+        let n = self.events.len();
+        if n == 0 {
+            self.color_groups = ColorGroups::new();
+            return;
+        }
+
+        let cg = color_greedy(n, |ev_idx| {
+            self.events[ev_idx].iter_agents().collect::<Vec<_>>()
+        });
+
+        let mut reordered: Vec<Event> = Vec::with_capacity(n);
+        let mut new_groups: Vec<Vec<usize>> = Vec::with_capacity(cg.groups.len());
+        let mut taken: Vec<Option<Event>> = self.events.drain(..).map(Some).collect();
+
+        for group in &cg.groups {
+            let mut new_indices: Vec<usize> = Vec::with_capacity(group.len());
+            for &old_idx in group {
+                let ev = taken[old_idx].take().expect("event already taken");
+                new_indices.push(reordered.len());
+                reordered.push(ev);
+            }
+            new_groups.push(new_indices);
+        }
+
+        self.events = reordered;
+        self.color_groups = ColorGroups { groups: new_groups };
+    }
+
     pub fn add_events<D: Drift<T>>(
         &mut self,
         composition: Vec<Vec<Vec<Index>>>,
@@ -212,6 +255,7 @@ impl<T: Time> TimeSlice<T> {
         self.events.extend(events);
 
         self.iteration(from, agents);
+        self.recompute_color_groups();
     }
 
     pub(crate) fn posteriors(&self) -> HashMap<Index, Gaussian> {
@@ -662,4 +706,67 @@ mod tests {
             epsilon = 1e-6
         );
     }
+
+    #[test]
+    fn time_slice_color_groups_reorders_events() {
+        // ev0: [a, b]; ev1: [c, d]; ev2: [a, c]
+        // Greedy coloring: ev0→c0, ev1→c0 (disjoint), ev2→c1 (overlaps both).
+        // After recompute_color_groups, physical order is [ev0, ev1, ev2]
+        // and groups == [[0, 1], [2]].
+        let mut index_map = KeyTable::new();
+
+        let a = index_map.get_or_create("a");
+        let b = index_map.get_or_create("b");
+        let c = index_map.get_or_create("c");
+        let d = index_map.get_or_create("d");
+
+        let mut agents: CompetitorStore<i64, ConstantDrift> = CompetitorStore::new();
+
+        for agent in [a, b, c, d] {
+            agents.insert(
+                agent,
+                Competitor {
+                    rating: Rating::new(
+                        Gaussian::from_ms(25.0, 25.0 / 3.0),
+                        25.0 / 6.0,
+                        ConstantDrift(25.0 / 300.0),
+                    ),
+                    ..Default::default()
+                },
+            );
+        }
+
+        let mut ts = TimeSlice::new(0i64, 0.0);
+
+        ts.add_events(
+            vec![
+                vec![vec![a], vec![b]],
+                vec![vec![c], vec![d]],
+                vec![vec![a], vec![c]],
+            ],
+            vec![vec![1.0, 0.0], vec![1.0, 0.0], vec![1.0, 0.0]],
+            vec![],
+            &agents,
+        );
+
+        assert_eq!(ts.color_groups.n_colors(), 2);
+        assert_eq!(ts.color_groups.groups[0], vec![0, 1]);
+        assert_eq!(ts.color_groups.groups[1], vec![2]);
+
+        assert_eq!(ts.color_groups.color_range(0), 0..2);
+        assert_eq!(ts.color_groups.color_range(1), 2..3);
+
+        // Events at positions 0 and 1 (color 0) must be disjoint — verify by
+        // checking that the agent sets of self.events[0] and self.events[1] do
+        // not include the agent at self.events[2].
+        let agents_in_ev2: Vec<Index> = ts.events[2].iter_agents().collect();
+        let agents_in_ev0: Vec<Index> = ts.events[0].iter_agents().collect();
+        let agents_in_ev1: Vec<Index> = ts.events[1].iter_agents().collect();
+        // ev0 and ev1 must be disjoint from each other (color-0 invariant).
+        assert!(agents_in_ev0.iter().all(|ag| !agents_in_ev1.contains(ag)));
+        // ev2 must share an agent with ev0 or ev1 (it needed its own color).
+        let ev2_overlaps_ev0 = agents_in_ev2.iter().any(|ag| agents_in_ev0.contains(ag));
+        let ev2_overlaps_ev1 = agents_in_ev2.iter().any(|ag| agents_in_ev1.contains(ag));
+        assert!(ev2_overlaps_ev0 || ev2_overlaps_ev1);
+    }
 }