function i = discreternd(p, varargin)
% DISCRETERND Random numbers from a discrete distribution
%
% Usage:
%   p = [.5 .25 .25];
%   i = discreternd(p) 
%      Generates a random number in {1,2,3} with 50 pct. probability of
%      being 1 and 25 pct. probability of being 2 or 3.
%
%   If p is a matrix, each row defines a discrete distribution, and a
%   random number is generated for each row.
%
%   i = discreternd(p, M, N, ...) 
%      Returns an size(p,1)-by-M-by-N-by-... array.
%
% Input:
%   p           Vector of (relative) probabilities. If p sums to one, each 
%               element in p is the probability of generating the 
%               corresponding integer. If p is not normalized it denotes 
%               the relative probability.
%   M, N, ...   Size of the returned array.
%
% Output:
%   i           Array of size M-by-N-by-... of random integers
%               between 1 and length(p).
%
% Copyright 2011, Mikkel N. Schmidt, Technical University of Denmark

if nargin>1
    n = varargin;
else
    n = {1};
end

if p==0
    i = 1;
else
    c = cumsum(bsxfun(@rdivide, p, sum(p,2)), 2);
    r = rand(size(p,1), n{:});
    i = nan(size(p,1), n{:});
    for k = 1:size(p,1)
        i(k,:) = sum(bsxfun(@gt, r(k,:), c(k,:)'),1)+1;
    end
end
