Chinese restaurant process

crp.jl

"""
Generate table assignments for 'n' customers,
according to a Chinese Restaurant Process with dispersion parameter 'α'
Returns an array of integer table assignments
"""
function crp(n, α)

    @assert n > 0
    table_assignments = zeros(Int, n)
    table_assignments[1] = 1 # first customer sits at table 1
    next_open_table = 2

    for i in 2:n
        if rand() < α/(α+i)
            # Customers sits at new table.
            table_assignments[i] = next_open_table
            next_open_table += 1
        else
            # Customer sits at an existing table.
            # He chooses which table to sit at by giving equal weight to each
            # customer already sitting at a table.
            which_table = rand(view(table_assignments,1:i-1))
            table_assignments[i] = which_table
        end
    end
    return table_assignments
end

using Plots
gr()
@time sample = crp(10^7, 10)
Plots.histogram(sample)