File Coverage

blib/lib/Statistics/WeightedSelection.pm

Criterion	Covered	Total	%
statement	115	118	97.4
branch	36	40	90.0
condition	7	11	63.6
subroutine	14	14	100.0
pod	7	7	100.0
total	179	190	94.2

package Statistics::WeightedSelection;
$Statistics::WeightedSelection::VERSION = '0.002';
use strict;
use warnings;
use Carp qw/croak cluck/;

use Storable qw/freeze/;

sub new {
    my ($class, %args) = @_;

    return bless {
        objects           => [],
        id_lookup         => {},
        replace_after_get => $args{replace_after_get} || 0,
    }, $class;
}

sub add {
    my ($self, %args) = @_;

    # a non-zero number which can include a decimal
    my $weight = $args{weight};

    # the scalar item that becomes a candidate for random selection
    my $object = $args{object};

    # an optional id, which can be used for later removal from pool
    my $id = defined $args{id} ? $args{id} : (ref $object ? freeze($object) : $object);

    unless ($weight && $object) {
        croak 'Calls to ' . __PACKAGE__ . "::add() must include an arg named 'object'"
            . " and a non-zero weight\n";
    }

    unless ($weight =~ /^\d+(\.\d*)?/) {
        croak 'Calls to ' . __PACKAGE__ . "::add()'s must include an arg named 'weight'"
            . " that must be a whole integer or number with decimal\n";
    }

    # in order to derive the starting_index (see the next stanza),
    # we need the last item in @$self
    my $last = $self->_get_last;

    # develops a structure that looks like the following, with each element being
    # [
    #     {
    #         starting_index => 0,
    #         length         => 40,
    #         object         => 'apple',
    #         id             => 'gwr3723',
    #     },
    #     {
    #         starting_index => 40,
    #         length         => 16,
    #         object         => 'plum',
    #         id             => 'avx9716',
    #     },
    #     {
    #         starting_index => 56,
    #         length         => 3.4
    #         object         => 'peach',
    #         id             => 'zzi1250',
    #     },
    #     {
    #         starting_index => 59.4,
    #         length         => 60,
    #         object         => 'mango',
    #         id             => 'umn2932',
    #     },
    # ]

    push @{$self->{objects}}, {
        # if this is the first object, our starting index is zero, and otherwise, we add the
        # new weight to the previous starting_index + weight for future binary search.
        starting_index => defined $last ? $last->{starting_index} + $last->{weight} : 0,
        weight         => $weight,
        object         => $object,
        id             => $id,
    };

    # I decided to use a hash with empty values for pointers to the index for a given
    # id or frozen object scalar, rather than using an arrayref, which might have been
    # more intuitive.  More than one added object can have the same id, and I wanted a
    # quick way to remove individual object pointers without having to iterate over an
    # array somehow.
    $self->{id_lookup}->{$id}->{$#{$self->{objects}}} = undef;

    return $#{$self->{objects}};
}

sub remove {
    my ($self, $id) = @_;

    unless (defined $id) {
        croak 'Calls to ' . __PACKAGE__ . "::remove() must include an id to remove\n";
    }

    $id = ref $id ? freeze($id) : $id;

    # delete the pointer altogether, as all of the objects with this id will be removed.
    my $indexes = delete $self->{id_lookup}->{$id};

    unless ($indexes && %{$indexes}) {
        # no need for this to be fatal, but it might indicate a bug in caller's code
        cluck "Key $id contains no associated indexes currently\n";
        return;
    }

    my @reverse_sorted_indexes = sort {$b <=> $a} keys %{$indexes};

    my @removed;
    for my $index (@reverse_sorted_indexes) {
        # remove all items that were pointed to that had this id
        # doing them in reverse, so that splicing out an earlier item won't
        # alter the indexes of items we want to splice out later.
        push @removed, splice(@{$self->{objects}}, $index, 1);
    }

    # we've altered our items, and need to fix our object starting indexes
    # and pointers for our keys.
    $self->_consolidate(reverse @reverse_sorted_indexes);

    # return the removed items
    return map {$_->{object}} @removed;
}

sub get {
    my ($self, $override_replacement) = @_;
    return unless @{ $self->{objects} };

    # when adding the starting_index and length together of the last item in @$self,
    #    (see the generated structure note in the add() method for more info), the
    #    max random number to generate is determined.
    my $last = $self->_get_last;
    my $random = rand($last->{starting_index} + $last->{weight});

    # binary search to quickly find the weighted index range.  the random number ($random)
    # is tested against the range of starting_index and (starting_index + length) to
    # determine if the number is lower or higher than the current arrayref until a match
    # is found.
    my $max = $#{ $self->{objects} };
    my $min = 0;
    my $index = 0;
    while ( $max >= $min ) {
        $index = int( ( $max + $min ) / 2 );
        my $current_object = $self->{objects}->[$index];

        if ( $random < $current_object->{starting_index} ) {
            $max = $index - 1;
        }
        elsif ( $random >= ($current_object->{starting_index} + $current_object->{weight}) ) {
            $min = $index + 1;
        }
        else {
            last;
        }
    }

    my $random_element;

    # don't run the removal logic in the else block if the user instructed us to replace
    # the object
    if ($self->replace_after_get) {
        $random_element = $self->{objects}->[$index];
    }
    else {
        # remove the element in question
        $random_element = splice(@{$self->{objects}}, $index, 1);

        # delete out the specific pointer to this item
        delete $self->{id_lookup}->{$random_element->{id}}->{$index};
        unless (keys %{$self->{id_lookup}->{$random_element->{id}}}) {
            delete $self->{id_lookup}->{$random_element->{id}};
        }

        # fix starting indexes and pointers for other objects.
        $self->_consolidate($index);
    }

    # return our selected object:
    return $random_element->{object};
}

sub replace_after_get {
    my ($self, $new_setting) = @_;
    $self->{replace_after_get} = $new_setting if defined $new_setting;
    return $self->{replace_after_get};
}

sub clear {
    my ($self) = @_;
    $self->{objects} = [];
    $self->{id_lookup} = {};
    return;
}

sub count {
    my ($self) = @_;
    return 0 if !@{ $self->{objects} };
    return $#{ $self->{objects} } + 1;
}

sub _dump {
    my ($self) = @_;
    # basically, don't return starting_index, as that is an internal implementation
    # detail, and doesn't need to be shown to the user
    return [map { {object => $_->{object}, weight => $_->{weight}, id => $_->{id}} }
        @{$self->{objects}}];
}

# internal method to fix id pointers and starting_indexes after a remove() or get() call
sub _consolidate {
    my $self = shift;

    # disregard all indexes greater than the current length of our objects property array
    my @removed_indexes = sort grep {$_ <= $#{ $self->{objects} }} @_;

    # separate our list of objects into segments bookended by our removed indexes
    # say we started with something like:
    #    0       1       2        3      4       5       6
    #    'alan', 'nate', 'brian', 'bob', 'nate', 'ryan', 'nate'
    #
    # and we called remove() on 'nate', then (1, 4, 6) would be passed to this routine, which
    # are array indexes, NOT our starting_indexes.  6 will be removed from consideration, as
    # it was at the end of the array.
    #
    # our first segment (former indexes 2 and 3) would be fixed as such:
    #     well, we need to change the following as a result:
    #         update 'brian''s starting index to be the starting index of 'alan' + its weight,
    #             and subtract 1 (as we removed 1 objects before this one) from its pointer in the id lookup.
    #         update 'bob''s starting index to be the starting index of 'brian' + its weight
    #             and subtract 1 (as we removed 1 objects before this one) from its pointer in the id lookup.
    #
    # our second segment (former index 5) would be fixed as such:
    #     well, we need to change the following as a result:
    #         update 'ryan''s starting index to be the starting index of 'bob' + its weight
    #             and subtract 2 (as we removed 2 objects before this one) from its pointer in the id lookup.

    for my $removed_index_index (0..$#removed_indexes) {
        # find our range bookends
        my $range_start = $removed_indexes[$removed_index_index];
        my $range_end   = $removed_index_index == $#removed_indexes
            ? $#{ $self->{objects} }
            : $removed_indexes[$removed_index_index + 1] - 1;

        # how many indexes were removed before our current segment, and thus the amount to subtract
        #     from the respective pointers
        my $to_subtract = @removed_indexes - $removed_index_index;

        my %ids_evaluated_for_range;
        for my $object_index ($range_start..$range_end) {
            my $object = $self->{objects}->[$object_index];

            # do all of the subtractions of pointer indexes at once for each segment.
            if (!$ids_evaluated_for_range{$object->{id}}++) {
                for my $index (
                    grep {$_ >= $range_start && ($removed_index_index == $#removed_indexes || $_ <= $range_end)}
                    keys %{$self->{id_lookup}->{$object->{id}}}
                ) {
                    delete $self->{id_lookup}->{$object->{id}}->{$index};
                    $self->{id_lookup}->{$object->{id}}->{$index - $to_subtract} = undef;
                }
            }

            # fix the starting indexes
            $object->{starting_index} = $object_index == 0
                ? 0
                : do {
                      my $previous_object = $self->{objects}->[$object_index - 1];
                      $previous_object->{starting_index} + $previous_object->{weight};
                  };
        }
    }

    return;
}


# simply a utility method to get the last item in the blessed array, as it contains all
# the information needed to add items and generate a random number.
sub _get_last {
    my ($self) = @_;
    return unless @{ $self->{objects} };
    return $self->{objects}->[$#{ $self->{objects} }];
}

1;

__END__