Я пытаюсь написать небольшую функцию для выполнения стратифицированной случайной выборки. То есть у меня есть вектор членства в группах для каждого элемента, и я хочу выбрать один элемент (индекс) для каждой группы. Таким образом, входными данными является количество желаемых элементов и членство в группе для каждого элемента. Результатом является список индексов.
Вот функция, которая у меня есть:
function stratified_sample(n::Int64, groups::Array{Int64})
# the output vector of indices
ind = zeros(Int64, n)
# first select n groups from the total set of possible groups
group_samp = sample(unique(groups), n, replace = false)
# cycle through the selected groups
for i in 1:n
# for each group, select one index whose group matches the current target group
ind[i] = sample([1:length(groups)...][groups.==group_samp[i]], 1, replace = false)[1]
end
# return the indices
return ind
end
Когда я запускаю этот код на относительно большом векторе, например, 1000 различных групп и 40000 общих записей , Я получаю
julia> groups = sample(1:1000, 40000, replace = true)
40000-element Array{Int64,1}:
221
431
222
421
714
108
751
259
⋮
199
558
317
848
271
358
julia> @time stratified_sample(5, groups)
0.022951 seconds (595.06 k allocations: 19.888 MiB)
5-element Array{Int64,1}:
11590
17057
17529
25103
20651
И для сравнения с нормальной случайной выборкой пяти элементов из возможных 40000:
julia> @time sample(1:40000, 5, replace = false)
0.000005 seconds (5 allocations: 608 bytes)
5-element Array{Int64,1}:
38959
5850
3283
19779
30063
Итак, мой код работает почти в 50 тысяч раз медленнее и использует в 33 тысячи раз больше памяти! Что, черт возьми, я сделал не так, и есть ли способ ускорить этот код? Я предполагаю, что реальное замедление происходит на этапе подмножества, т.е. [1:length(groups)...][groups.==group_samp[i]], но я не могу найти лучшего решения.
Я бесконечно искал эту функцию в стандартных пакетах Julia, но безуспешно.
Есть предложения?
EDIT: Мне удалось значительно ускорить его, просто взяв случайную выборку и проверив, есть ли она удовлетворяет требованию, чтобы было выбрано n уникальных групп:
function stratified_sample_random(n::Int64, groups::Array{Int64}, group_probs::Array{Float32})
ind = zeros(Int64, n)
my_samp = []
while true
my_samp = wsample(1:length(groups), group_probs, n, replace = false)
if length(unique(groups[my_samp])) == n
break
end
end
return my_samp
end
Здесь group_probs - это просто вектор вероятностей выборки, где элементы каждой группы имеют общую вероятность 1 / с, где s - это количество элементов в этой группе. Например, если groups = [1,1,1,1,2,3,3], соответствующие вероятности будут group_probs = [0.25, 0.25, 0.25, 0.25, 1, 0.5, 0.5]. Это помогает ускорить выборку, сводя к минимуму вероятность выбора нескольких элементов одной группы. В целом это работает довольно хорошо:
@time stratified_sample_random(5, groups, group_probs)
0.000122 seconds (14 allocations: 1.328 KiB)
5-element Array{Int64,1}:
32209
10184
30892
4861
30300
После небольшого экспериментирования взвешенная выборка по вероятности не обязательно быстрее стандартной выборки (), но это зависит от того, сколько уникальных групп и что требуется n значение есть.
Конечно, нет гарантии, что эта функция будет случайным образом выбирать уникальный набор объектов, и она может циклически повторяться бесконечно. Моя мысль состоит в том, чтобы добавить счетчик к while l oop, и если он попробует что-то вроде 10000 раз безуспешно, он вызовет исходную функцию stratified_sample, которую я представил, чтобы гарантировать, что она вернет уникальный результат. Мне не нравится это решение, и должен быть более элегантный и экономный подход, но это определенно улучшение.