Group Sums Without a GroupBy

Inside a training loop there is no DataFrame. Per-user, per-slate and per-token aggregates are integer-indexed arrays, and this is how they are computed.

Implement segment_stats(values, group_ids, n_groups). Given a float array values of length n, an integer array group_ids of the same length with every id in [0, n_groups), return a tuple (sums, means) where both are float arrays of length n_groups.

values    = [1.0, 2.0, 3.0, 4.0]
group_ids = [0,   0,   1,   1  ]
n_groups  = 2
sums      = [3.0, 7.0]
means     = [1.5, 3.5]

A group that no row belongs to has a sum of 0.0 and a mean of 0.0, not nan. That is a deliberate contract: this array feeds a model, and one nan in a feature column poisons every gradient computed from the batch.

Do it without any Python loops.