!YS-v0

defn main(n=5000):
  :: https://rosettacode.org/wiki/100_prisoners

  random-successes optimal-successes run-count =:
    simulate-n-runs(n)
      .slice(q(random-successes optimal-successes run-count))

  say: "Probability of survival with random search: $F(random-successes / run-count)"

  say: "Probability of survival with ordered search: $F(optimal-successes / run-count)"

defn simulate-n-runs(n):
  :: Simulate n runs of the 100 prisoner problem and returns a success count
     for each search method.

  loop random-successes 0, optimal-successes 0, run-count 0:
    # If we've done the loop n times
    if n == run-count:
      # return results
      then::
        random-successes  :: random-successes
        optimal-successes :: optimal-successes
        run-count         :: run-count

      # Otherwise, run for another batch of prisoners
      else:
        next-result =: simulate-100-prisoners()
        recur:
          (random-successes + next-result.random)
          (optimal-successes + next-result.optimal)
          run-count.++

defn simulate-100-prisoners():
  :: Simulates all prisoners searching the same drawers by both strategies,
     returns map showing whether each was successful.

  # Create 100 drawers with randomly ordered prisoner numbers:
  drawers =: random-drawers()
  hash-map:
    :random   try-luck(drawers search-50-random-drawers)
    :optimal  try-luck(drawers search-50-optimal-drawers)

defn try-luck(drawers drawer-searching-function):
  :: Returns 1 if all prisoners find their number otherwise 0.

  loop prisoners range(100):
    if prisoners.?:
      if prisoners.0.drawer-searching-function(drawers):
        recur: rest(prisoners)
        else: 0
      else: 1

defn search-50-optimal-drawers(prisoner-number drawers):
  :: Open 50 drawers according to the agreed strategy, returning true if
     prisoner's number was found.

  loop next-drawer prisoner-number, drawers-opened 0:
    when drawers-opened != 50:
      result =: drawers.$next-drawer
      result == prisoner-number ||:
        recur: result drawers-opened.++

defn search-50-random-drawers(prisoner-number drawers):
  :: Select 50 random drawers and return true if the prisoner's number was
     found.

  drawers:
    .shuffle()
    .take(50)
    .filter(eq(prisoner-number))
    .count()
    .eq(1)

defn random-drawers():
  :: Returns a list of shuffled numbers.

  shuffle: range(100)