numpy.random.choiceみたいなやつを実現するのに結構苦労した。

Rustで乱数を扱うにはrandというパッケージを使うのがいいらしい。

rand - Rust

しかし、ドキュメントを読む感じだと重み付きサンプリングをやる方法が書いていなくて苦労した。色々調べた結果

How can I randomly select one element from a vector or array?

WeightedIndexというのがそんな感じの機能らしい。例としてあげられていたサンプルコードはこんな感じ

use rand::prelude::*;
use rand::distributions::WeightedIndex;

let choices = ['a', 'b', 'c'];
let weights = [2, 1, 1];
let dist = WeightedIndex::new(&weights).unwrap();
let mut rng = thread_rng();
for _ in 0..100 {
	// 50% chance to pring 'a', 25% change to pring 'b', 25% change to print 'c'
	println!("{}", choices[dist.sample(&mut rng)]);
}

let items = [('a', 0), ('b', 3), ('c', 7)];
let dist2 = WeightedIndex::new(items.iter().map(|item| item.1)).unwrap();
for _ in 0..100 {
	// 0% chance to print 'a', 30% chance to print 'b', 70% chance to print 'c'
  println!("{}", items[dist2.sample(&mut rng)].0);
}

今、やりたいのはDataFrameのlidarというカラムの値の出現回数を数え上げ、それを確率値に直し、確率値に従って一つ値をサンプルする、というものだ(「詳解」の本でやっている)。

このうち最初の二つはそこまで苦労しなかった。

let lidar = df.column("lidar").unwrap();
let mut value_counts = lidar.value_counts().unwrap();  // 返り値はDataFrame

let counts = value_counts.column("counts").unwrap();
let mut probs: Series = counts
	.i64()
	.unwrap()
	.into_iter()
	.map(|opt_x| match opt_x {
		Some(x) => Some((x as f64) / (lidar.len() as f64)),
		None => None
	})
	.collect();

probs.rename("probs");
value_counts = value_counts.with_column(probs).unwrap();

この結果次のようなDataFrameが得られる

このlidarのカラムをprobsに従ってサンプリングするのが目標だ。やり方としてはprobsを元にWeightedIndexを作り、それを使ってlidarからサンプリングを行う。ただし、素のSeriesはNoneが含まれる可能性があるため、おそらくWeightedIndexを作れないのでまず使いやすくするためにVectorに直してやる。

let lidar_vec = value_counts
	.column("lidar")
	.unwrap()
	.i64()
	.unwrap()
	.into_iter()
	.map(|opt_x| opt_x.unwrap())
	.collect::<Vec<i64>>();

let probs_vec = value_counts
	.column("probs")
	.unwrap()
	.f64()
	.unwrap()
	.into_iter()
	.map(|opt_x| opt_x.unwrap())
	.collect::<Vec<f64>>();

その上で以下のようにするとサンプリングができる

let dist = WeightedIndex::new(&probs_vec).unwrap();
let mut rng = thread_rng();
let sampled = lidar_vec[dist.sample(&mut rng)];