mirror of
https://github.com/pocketpy/pocketpy
synced 2025-10-20 11:30:18 +00:00
101 lines
4.1 KiB
C++
101 lines
4.1 KiB
C++
#include "pocketpy/random.h"
|
|
|
|
namespace pkpy{
|
|
|
|
struct Random{
|
|
PY_CLASS(Random, random, Random)
|
|
std::mt19937 gen;
|
|
|
|
Random(){
|
|
gen.seed(std::chrono::high_resolution_clock::now().time_since_epoch().count());
|
|
}
|
|
|
|
static void _register(VM* vm, PyObject* mod, PyObject* type){
|
|
vm->bind_default_constructor<Random>(type);
|
|
|
|
vm->bind_method<1>(type, "seed", [](VM* vm, ArgsView args) {
|
|
Random& self = _CAST(Random&, args[0]);
|
|
self.gen.seed(CAST(i64, args[1]));
|
|
return vm->None;
|
|
});
|
|
|
|
vm->bind_method<2>(type, "randint", [](VM* vm, ArgsView args) {
|
|
Random& self = _CAST(Random&, args[0]);
|
|
i64 a = CAST(i64, args[1]);
|
|
i64 b = CAST(i64, args[2]);
|
|
if (a > b) vm->ValueError("randint(a, b): a must be less than or equal to b");
|
|
std::uniform_int_distribution<i64> dis(a, b);
|
|
return VAR(dis(self.gen));
|
|
});
|
|
|
|
vm->bind_method<0>(type, "random", [](VM* vm, ArgsView args) {
|
|
Random& self = _CAST(Random&, args[0]);
|
|
std::uniform_real_distribution<f64> dis(0.0, 1.0);
|
|
return VAR(dis(self.gen));
|
|
});
|
|
|
|
vm->bind_method<2>(type, "uniform", [](VM* vm, ArgsView args) {
|
|
Random& self = _CAST(Random&, args[0]);
|
|
f64 a = CAST(f64, args[1]);
|
|
f64 b = CAST(f64, args[2]);
|
|
std::uniform_real_distribution<f64> dis(a, b);
|
|
return VAR(dis(self.gen));
|
|
});
|
|
|
|
vm->bind_method<1>(type, "shuffle", [](VM* vm, ArgsView args) {
|
|
Random& self = _CAST(Random&, args[0]);
|
|
List& L = CAST(List&, args[1]);
|
|
std::shuffle(L.begin(), L.end(), self.gen);
|
|
return vm->None;
|
|
});
|
|
|
|
vm->bind_method<1>(type, "choice", [](VM* vm, ArgsView args) {
|
|
Random& self = _CAST(Random&, args[0]);
|
|
auto [data, size] = vm->_cast_array(args[1]);
|
|
if(size == 0) vm->IndexError("cannot choose from an empty sequence");
|
|
std::uniform_int_distribution<i64> dis(0, size - 1);
|
|
return data[dis(self.gen)];
|
|
});
|
|
|
|
vm->bind(type, "choices(self, population, weights=None, k=1)", [](VM* vm, ArgsView args) {
|
|
Random& self = _CAST(Random&, args[0]);
|
|
auto [data, size] = vm->_cast_array(args[1]);
|
|
if(size == 0) vm->IndexError("cannot choose from an empty sequence");
|
|
std::vector<f64> cum_weights(size);
|
|
if(args[2] == vm->None){
|
|
for(int i = 0; i < size; i++) cum_weights[i] = i + 1;
|
|
}else{
|
|
auto [weights, weights_size] = vm->_cast_array(args[2]);
|
|
if(weights_size != size) vm->ValueError(_S("len(weights) != ", size));
|
|
cum_weights[0] = CAST(f64, weights[0]);
|
|
for(int i = 1; i < size; i++){
|
|
cum_weights[i] = cum_weights[i - 1] + CAST(f64, weights[i]);
|
|
}
|
|
}
|
|
if(cum_weights[size - 1] <= 0) vm->ValueError("total of weights must be greater than zero");
|
|
int k = CAST(i64, args[3]);
|
|
List result(k);
|
|
for(int i = 0; i < k; i++){
|
|
f64 r = std::uniform_real_distribution<f64>(0.0, cum_weights[size - 1])(self.gen);
|
|
int idx = std::lower_bound(cum_weights.begin(), cum_weights.end(), r) - cum_weights.begin();
|
|
result[i] = data[idx];
|
|
}
|
|
return VAR(std::move(result));
|
|
});
|
|
}
|
|
};
|
|
|
|
void add_module_random(VM* vm){
|
|
PyObject* mod = vm->new_module("random");
|
|
Random::register_class(vm, mod);
|
|
PyObject* instance = vm->heap.gcnew<Random>(Random::_type(vm));
|
|
mod->attr().set("seed", vm->getattr(instance, "seed"));
|
|
mod->attr().set("random", vm->getattr(instance, "random"));
|
|
mod->attr().set("uniform", vm->getattr(instance, "uniform"));
|
|
mod->attr().set("randint", vm->getattr(instance, "randint"));
|
|
mod->attr().set("shuffle", vm->getattr(instance, "shuffle"));
|
|
mod->attr().set("choice", vm->getattr(instance, "choice"));
|
|
mod->attr().set("choices", vm->getattr(instance, "choices"));
|
|
}
|
|
|
|
} // namespace pkpy
|