├── .gitignore ├── Gemfile ├── LICENSE ├── README.md ├── Rakefile ├── benchmarks └── pickup.rb ├── lib ├── pickup.rb └── pickup │ ├── circle_iterator.rb │ ├── mapped_list.rb │ └── version.rb ├── pickup.gemspec └── spec ├── pickup └── pickup_spec.rb └── spec_helper.rb /.gitignore: -------------------------------------------------------------------------------- 1 | *.gem 2 | *.rbc 3 | .bundle 4 | .config 5 | .yardoc 6 | Gemfile.lock 7 | InstalledFiles 8 | _yardoc 9 | coverage 10 | doc/ 11 | lib/bundler/man 12 | pkg 13 | rdoc 14 | spec/reports 15 | test/tmp 16 | test/version_tmp 17 | tmp 18 | -------------------------------------------------------------------------------- /Gemfile: -------------------------------------------------------------------------------- 1 | source 'https://rubygems.org' 2 | 3 | # Specify your gem's dependencies in pickup.gemspec 4 | gemspec 5 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (c) 2012 fl00r 2 | 3 | MIT License 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining 6 | a copy of this software and associated documentation files (the 7 | "Software"), to deal in the Software without restriction, including 8 | without limitation the rights to use, copy, modify, merge, publish, 9 | distribute, sublicense, and/or sell copies of the Software, and to 10 | permit persons to whom the Software is furnished to do so, subject to 11 | the following conditions: 12 | 13 | The above copyright notice and this permission notice shall be 14 | included in all copies or substantial portions of the Software. 15 | 16 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 17 | EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 18 | MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 19 | NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 20 | LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 21 | OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 22 | WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Pickup 2 | 3 | Pickup helps you to pick an item from a collection by its weight or probability. 4 | 5 | ## Installation 6 | 7 | Add this line to your application's Gemfile: 8 | 9 | gem 'pickup' 10 | 11 | And then execute: 12 | 13 | bundle 14 | 15 | Or install it yourself as: 16 | 17 | gem install pickup 18 | 19 | ## Usage 20 | 21 | For example, we have got a pond with fish. 22 | 23 | ```ruby 24 | pond = { 25 | "selmon" => 1, 26 | "carp" => 4, 27 | "crucian" => 3, 28 | "herring" => 6, 29 | "sturgeon" => 8, 30 | "gudgeon" => 10, 31 | "minnow" => 20 32 | } 33 | ``` 34 | 35 | Values are a chance (probability) to get this fish. 36 | 37 | So we should create our pickup. 38 | 39 | ```ruby 40 | pickup = Pickup.new(pond) 41 | pickup.pick(3) 42 | #=> [ "gudgeon", "minnow", "minnow" ] 43 | ``` 44 | 45 | Look, we've just catched few minnows! To get selmon we need some more tries ;) 46 | 47 | ### Custom distribution function 48 | 49 | Ok. What if our probability is not a linear function. We can create our pickup with a function: 50 | 51 | ```ruby 52 | pickup = Pickup.new(pond){ |v| v**2 } 53 | pickup.pick(3) 54 | #=> ["carp", "selmon", "crucian"] 55 | ``` 56 | 57 | Wow, good catch! 58 | 59 | Also you can change our "function" on the fly. Let's make square function: 60 | 61 | ```ruby 62 | pickup = Pickup.new(pond) 63 | pickup.pick_func = Proc.new{ |v| v**2 } 64 | ``` 65 | 66 | Or you can pass a block as a probability function wich will be applicable only to current operation 67 | 68 | ```ruby 69 | pickup = Pickup.new(pond) 70 | pickup.pick{ |v| Math.sin(v) } # same as pickup.pick(1){ ... } 71 | #=> "selmon" 72 | pickup.pick 73 | #=> "minnow" 74 | ``` 75 | 76 | In case of `f(weight)=weight^10` most possible result will be "minnow", because `20^10` is `2^10` more possible then "gudgeon" 77 | 78 | ```ruby 79 | pickup = Pickup.new(pond) 80 | pickup.pick(10){ |v| v**10 } 81 | #=> ["minnow", "minnow", "minnow", "minnow", "minnow", "minnow", "minnow", "minnow", "minnow", "minnow"] 82 | ``` 83 | 84 | Or you can use reverse probability: 85 | 86 | ```ruby 87 | pickup = Pickup.new(pond) 88 | pickup.pick(10){ |v| v**(-10) } 89 | #=> ["selmon", "selmon", "selmon", "selmon", "crucian", "selmon", "selmon", "selmon", "selmon", "selmon"] 90 | ``` 91 | 92 | ### Random uniq pick 93 | 94 | Also we can pick random uniq items from the list 95 | 96 | ```ruby 97 | pickup = Pickup.new(pond, uniq: true) 98 | pickup.pick(3) 99 | #=> [ "gudgeon", "herring", "minnow" ] 100 | pickup.pick 101 | #=> "herring" 102 | pickup.pick 103 | #=> "gudgeon" 104 | pickup.pick 105 | #=> "sturgeon" 106 | ``` 107 | 108 | ### Custom key and weight selection functions 109 | 110 | We can use more complex collections by defining our own key and weight selectors: 111 | 112 | ```ruby 113 | require "ostruct" 114 | 115 | pond_ostruct = [ 116 | OpenStruct.new(key: "sel", name: "selmon", weight: 1), 117 | OpenStruct.new(key: "car", name: "carp", weight: 4), 118 | OpenStruct.new(key: "cru", name: "crucian", weight: 3), 119 | OpenStruct.new(key: "her", name: "herring", weight: 6), 120 | OpenStruct.new(key: "stu", name: "sturgeon", weight: 8), 121 | OpenStruct.new(key: "gud", name: "gudgeon", weight: 10), 122 | OpenStruct.new(key: "min", name: "minnow", weight: 20) 123 | ] 124 | 125 | key_func = Proc.new{ |item| item.key } 126 | weight_func = Proc.new{ |item| item.weight } 127 | pickup = Pickup.new(pond_ostruct, key_func: key_func, weight_func: weight_func) 128 | 129 | # Symbol values for funcs will be converted into Procs: 130 | pickup = Pickup.new(pond_ostruct, key_func: :key, weight_func: :weight) 131 | 132 | pickup.pick 133 | #=> "gud" 134 | 135 | name_func = Proc.new{ |item| item.name } 136 | pickup.pick(1, key_func: name_func) 137 | #=> "gudgeon" 138 | ``` 139 | 140 | ## Contributing 141 | 142 | 1. Fork it 143 | 2. Create your feature branch (`git checkout -b my-new-feature`) 144 | 3. Commit your changes (`git commit -am 'Added some feature'`) 145 | 4. Push to the branch (`git push origin my-new-feature`) 146 | 5. Create new Pull Request 147 | -------------------------------------------------------------------------------- /Rakefile: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env rake 2 | require "bundler/gem_tasks" 3 | require 'rake/testtask' 4 | 5 | task :default => :spec 6 | 7 | desc 'Tests' 8 | Rake::TestTask.new(:spec) do |t| 9 | t.libs << 'spec' 10 | t.pattern = 'spec/**/*_spec.rb' 11 | t.verbose = false 12 | end -------------------------------------------------------------------------------- /benchmarks/pickup.rb: -------------------------------------------------------------------------------- 1 | require 'benchmark' 2 | require './lib/pickup' 3 | 4 | def simple_hash 5 | @simple ||= begin 6 | hash = {} 7 | (1..30).to_a.each do |i| 8 | hash["item_#{i}"] = rand(30) 9 | end 10 | hash 11 | end 12 | end 13 | 14 | def big_hash 15 | @big ||= begin 16 | hash = {} 17 | (1..5000).to_a.each do |i| 18 | hash["item_#{i}"] = rand(1000) 19 | end 20 | hash 21 | end 22 | end 23 | 24 | def big_weights_hash 25 | @big_weights ||= begin 26 | hash = {} 27 | (1..5000).to_a.each do |i| 28 | hash["item_#{i}"] = rand(100_000) + 10_000 29 | end 30 | hash 31 | end 32 | end 33 | 34 | def simple(uniq=false) 35 | pickup = Pickup.new(simple_hash, uniq: uniq) 36 | pickup.pick(10) 37 | end 38 | 39 | def big(uniq=false) 40 | pickup = Pickup.new(big_hash, uniq: uniq) 41 | pickup.pick(100) 42 | end 43 | 44 | def big_weights(uniq=false) 45 | pickup = Pickup.new(big_weights_hash, uniq: uniq) 46 | pickup.pick(100) 47 | end 48 | 49 | n = 500 50 | 51 | Benchmark.bm do |x| 52 | x.report("simple: ") do 53 | n.times{ simple } 54 | end 55 | x.report("simple uniq: ") do 56 | n.times{ simple(true) } 57 | end 58 | x.report("big: ") do 59 | n.times{ big } 60 | end 61 | x.report("big uniq: ") do 62 | n.times{ big(true) } 63 | end 64 | x.report("big weights: ") do 65 | n.times{ big_weights } 66 | end 67 | x.report("big weights uniq: ") do 68 | n.times{ big_weights(true) } 69 | end 70 | end -------------------------------------------------------------------------------- /lib/pickup.rb: -------------------------------------------------------------------------------- 1 | require 'pickup/version' 2 | 3 | class Pickup 4 | require 'pickup/circle_iterator' 5 | require 'pickup/mapped_list' 6 | 7 | attr_reader :list, :uniq 8 | attr_writer :pick_func, :key_func, :weight_func 9 | 10 | def initialize(list, opts = {}, &block) 11 | @list = list 12 | @uniq = opts[:uniq] || false 13 | @pick_func = block if block_given? 14 | @key_func = Pickup.func_opt(opts[:key_func]) 15 | @weight_func = Pickup.func_opt(opts[:weight_func]) 16 | end 17 | 18 | def pick(count = 1, opts = {}, &block) 19 | func = block || pick_func 20 | key_func = Pickup.func_opt(opts[:key_func]) || @key_func 21 | weight_func = Pickup.func_opt(opts[:weight_func]) || @weight_func 22 | mlist = MappedList.new(list, func, uniq: uniq, key_func: key_func, weight_func: weight_func) 23 | result = mlist.random(count) 24 | count == 1 ? result.first : result 25 | end 26 | 27 | def pick_func 28 | @pick_func ||= proc { |val| val } 29 | end 30 | 31 | def self.func_opt(opt) 32 | opt.is_a?(Symbol) ? opt.to_proc : opt 33 | end 34 | end 35 | -------------------------------------------------------------------------------- /lib/pickup/circle_iterator.rb: -------------------------------------------------------------------------------- 1 | class Pickup 2 | class CircleIterator 3 | attr_reader :func, :obj, :max 4 | 5 | def initialize(obj, func, max, opts = {}) 6 | @obj = obj.dup 7 | @func = func 8 | @max = max 9 | @key_func = Pickup.func_opt(opts[:key_func]) || key_func 10 | @weight_func = Pickup.func_opt(opts[:weight_func]) || weight_func 11 | end 12 | 13 | def key_func 14 | @key_func ||= proc { |item| item[0] } 15 | end 16 | 17 | def weight_func 18 | @weight_func ||= proc { |item| item[1] } 19 | end 20 | 21 | def each 22 | until obj.empty? 23 | start = 0 24 | 25 | obj.each do |item| 26 | key = key_func.call(item) 27 | weight = weight_func.call(item) 28 | 29 | val = func.call(weight) 30 | start += val 31 | 32 | if yield([key, start, max]) 33 | obj.delete(key) 34 | @max -= val 35 | end 36 | end 37 | end 38 | end 39 | end 40 | end 41 | -------------------------------------------------------------------------------- /lib/pickup/mapped_list.rb: -------------------------------------------------------------------------------- 1 | class Pickup 2 | class MappedList 3 | attr_reader :list, :func, :uniq 4 | 5 | BOOLEAN_DEPRECATION ||= '[DEPRECATED] Passing uniq as a boolean to ' \ 6 | "MappedList's initialize method is deprecated. Please use the opts hash " \ 7 | 'instead.'.freeze 8 | 9 | def initialize(list, func, opts = nil) 10 | if opts.is_a?(Hash) 11 | @key_func = Pickup.func_opt(opts[:key_func]) || key_func 12 | @weight_func = Pickup.func_opt(opts[:weight_func]) || weight_func 13 | @uniq = opts[:uniq] || false 14 | else 15 | # If opts is explicitly provided as a boolean, show the warning. 16 | warn BOOLEAN_DEPRECATION if [true, false].include?(opts) 17 | 18 | @uniq = opts || false 19 | end 20 | 21 | @func = func 22 | @list = list 23 | @current_state = 0 24 | end 25 | 26 | def key_func 27 | @key_func ||= proc { |item| item[0] } 28 | end 29 | 30 | def weight_func 31 | @weight_func ||= proc { |item| item[1] } 32 | end 33 | 34 | def max 35 | @max ||= list.sum { |item| func.call(weight_func.call(item)) } 36 | end 37 | 38 | def each(*) 39 | CircleIterator.new( 40 | @list, func, max, key_func: key_func, weight_func: weight_func 41 | ).each do |item| 42 | if uniq 43 | true if yield item 44 | else 45 | nil while yield(item) 46 | end 47 | end 48 | end 49 | 50 | def random(count) 51 | if uniq && list.size < count 52 | raise 'List is shorter than count of items you want to get' 53 | end 54 | 55 | nums = count.times.map { rand(max) }.sort 56 | return [] if max.zero? 57 | 58 | get_random_items(nums) 59 | end 60 | 61 | def get_random_items(nums) 62 | current_num = nums.shift 63 | items = [] 64 | each do |item, counter, mx| 65 | break unless current_num 66 | next unless counter % (mx + 1) > current_num % mx 67 | 68 | items << item 69 | current_num = nums.shift 70 | end 71 | items 72 | end 73 | end 74 | end 75 | -------------------------------------------------------------------------------- /lib/pickup/version.rb: -------------------------------------------------------------------------------- 1 | class Pickup 2 | VERSION = "0.0.12" 3 | end 4 | -------------------------------------------------------------------------------- /pickup.gemspec: -------------------------------------------------------------------------------- 1 | # -*- encoding: utf-8 -*- 2 | require File.expand_path('../lib/pickup/version', __FILE__) 3 | 4 | Gem::Specification.new do |gem| 5 | gem.authors = ["fl00r"] 6 | gem.email = ["fl00r@yandex.ru"] 7 | gem.description = %q{Pickup helps you to pick item from collection by it's weight/probability} 8 | gem.summary = %q{Pickup helps you to pick item from collection by it's weight/probability} 9 | gem.homepage = "" 10 | 11 | gem.files = `git ls-files`.split($\) 12 | gem.executables = gem.files.grep(%r{^bin/}).map{ |f| File.basename(f) } 13 | gem.test_files = gem.files.grep(%r{^(test|spec|features)/}) 14 | gem.name = "pickup" 15 | gem.require_paths = ["lib"] 16 | gem.version = Pickup::VERSION 17 | gem.required_ruby_version = ">= 1.9" 18 | end 19 | -------------------------------------------------------------------------------- /spec/pickup/pickup_spec.rb: -------------------------------------------------------------------------------- 1 | # encoding: utf-8 2 | require 'spec_helper' 3 | require 'ostruct' 4 | require 'stringio' 5 | 6 | describe Pickup do 7 | 8 | before do 9 | @list = { 10 | "selmon" => 1, # 1 11 | "carp" => 4, # 5 12 | "crucian" => 3, # 8 13 | "herring" => 6, # 14 14 | "sturgeon" => 8, # 22 15 | "gudgeon" => 10, # 32 16 | "minnow" => 20 # 52 17 | } 18 | 19 | @struct_list = @list.map{ |key, weight| OpenStruct.new(key: key, weight: weight) } 20 | 21 | @func = Proc.new{ |a| a } 22 | @pickup = Pickup.new(@list) 23 | @pickup2 = Pickup.new(@list, uniq: true) 24 | 25 | @key_func = Proc.new{ |item| item.key } 26 | @weight_func = Proc.new{ |item| item.weight } 27 | @pickup3 = Pickup.new(@struct_list, key_func: @key_func, weight_func: @weight_func) 28 | end 29 | 30 | it "should pick correct amount of items" do 31 | @pickup.pick(2).size.must_equal 2 32 | @pickup.pick(10).size.must_equal 10 33 | end 34 | 35 | describe Pickup::MappedList do 36 | before do 37 | @ml = Pickup::MappedList.new(@list, @func, uniq: true) 38 | @ml2 = Pickup::MappedList.new(@list, @func) 39 | @ml3 = Pickup::MappedList.new(@struct_list, @func, key_func: @key_func, weight_func: @weight_func) 40 | @ml4 = Pickup::MappedList.new(@struct_list, @func, uniq: true, key_func: @key_func, weight_func: @weight_func) 41 | end 42 | 43 | it "deprecated warning on initialization if uniq is passed directly as a boolean" do 44 | # Swap in a fake IO object for $stderr. 45 | orig_stderr = $stderr 46 | $stderr = StringIO.new 47 | 48 | Pickup::MappedList.new(@list, @func, true) 49 | 50 | # Inspect the fake IO object for the deprecated warning. 51 | $stderr.rewind 52 | $stderr.string.chomp.must_equal("[DEPRECATED] Passing uniq as a boolean to MappedList's initialize method is deprecated. Please use the opts hash instead.") 53 | 54 | # Restore the original stderr. 55 | $stderr = orig_stderr 56 | end 57 | 58 | it "should return selmon and then carp and then crucian for uniq pickup" do 59 | @ml.get_random_items([0, 0, 0]).must_equal ["selmon", "carp", "crucian"] 60 | end 61 | 62 | it "should return selmon 3 times for non-uniq pickup" do 63 | @ml2.get_random_items([0]).first.must_equal "selmon" 64 | @ml2.get_random_items([0]).first.must_equal "selmon" 65 | @ml2.get_random_items([0]).first.must_equal "selmon" 66 | end 67 | 68 | it "should return crucian 3 times for uniq pickup" do 69 | @ml2.get_random_items([7, 7, 7]).must_equal ["crucian", "crucian", "crucian"] 70 | end 71 | 72 | it "should return item from the beginning after end of list for uniq pickup" do 73 | @ml.get_random_items([20, 20, 20, 20]).must_equal ["sturgeon", "gudgeon", "minnow", "crucian"] 74 | end 75 | 76 | it "should return right max" do 77 | @ml.max.must_equal 52 78 | end 79 | 80 | it "should return selmon 4 times for non-uniq pickup (using custom weight function)" do 81 | 4.times{ @ml3.get_random_items([0]).first.must_equal "selmon" } 82 | end 83 | 84 | it "should return right max (using custom weight function)" do 85 | @ml3.max.must_equal 52 86 | end 87 | 88 | it "should return selmon and then carp and then crucian for uniq pickup (using custom weight function)" do 89 | @ml4.get_random_items([0, 0, 0]).must_equal ["selmon", "carp", "crucian"] 90 | end 91 | end 92 | 93 | it "should take 7 different fish" do 94 | items = @pickup2.pick(7) 95 | items.uniq.size.must_equal 7 96 | end 97 | 98 | it "should raise an exception" do 99 | proc{ items = @pickup2.pick(8) }.must_raise RuntimeError 100 | end 101 | 102 | it "should return a list of items including the heaviest item (but not always - sometimes it will fail)" do 103 | items = @pickup2.pick(2){ |v| v**20 } 104 | (items.include? "minnow").must_equal true 105 | end 106 | 107 | it "should return a list of items including the lightest item (but not always - sometimes it will fail)" do 108 | items = @pickup2.pick(2){ |v| v**(-20) } 109 | (items.include? "selmon").must_equal true 110 | end 111 | 112 | it "should pick correct amount of items (using custom weight function)" do 113 | @pickup3.pick(4).size.must_equal 4 114 | @pickup3.pick(12).size.must_equal 12 115 | end 116 | 117 | it "should take 5 fish (using custom weight function)" do 118 | @pickup3.pick(5, key_func: @key_func, weight_func: @weight_func).size.must_equal 5 119 | end 120 | 121 | let(:list) { { "foo" => 0 } } 122 | let(:pickup) { Pickup.new(list) } 123 | 124 | it "should not devide by zero" do 125 | pickup.pick(1) 126 | end 127 | end -------------------------------------------------------------------------------- /spec/spec_helper.rb: -------------------------------------------------------------------------------- 1 | require 'pickup' 2 | require 'minitest/spec' 3 | require 'minitest/autorun' --------------------------------------------------------------------------------