Brute-force for survivor game v1.11

#/usr/bin/env pypy

from datetime import datetime
from fractions import Fraction
import operator
import sys

def eliminate(n, p_map):
  survivors = set(range(1, n+1))
  for i in range(n, 0, -1):
    if i in survivors:
      survivors.discard(p_map[i])
  return survivors

def average_eval_prob_maps(n, evals):
  l = len(evals)
  return {i: Fraction(1, l) * sum([e[2][i] for e in evals]) for i in range(1, n+1)}

# return (probability whether current player is alive, # survivors, map of k -> probability that k is alive in this path))
def evaluate_dfs(n, cur_path):
  k = len(cur_path) # how deep we are in the decision tree
  # print('cur_path =', cur_path)
  if k == n: # time to evaluate
    survivors = eliminate(n, cur_path)
    prob_map = {i: Fraction(1, 1) if i in survivors else Fraction(0,1) for i in range(1, n+1)}
    return (prob_map[k], -len(survivors), prob_map) # want to sort the opposite direction by # survivors
  else:
    new_paths = [{**cur_path, **{k+1: i}} for i in range(1, n+1)]
    new_path_evals = [evaluate_dfs(n, path) for path in new_paths]
    # print("NEW PATHS")
    # print(list(zip(new_paths, new_path_evals)))
    new_path_evals.sort(key=operator.itemgetter(0, 1))
    # now sorted by [survived, # survivors]
    # get all paths tied with the best path, have to choose randomly between them
    best_path_eval = new_path_evals[-1]
    best_path_evals = [best_path_eval]
    for path_eval in reversed(new_path_evals[:-1]):
      if path_eval[0] == best_path_eval[0] and path_eval[1] == best_path_eval[1]:
        best_path_evals.append(path_eval)

    # print("BEST_PATH_EVALS=", best_path_evals)

    prob_map = average_eval_prob_maps(n, best_path_evals)
    if len(cur_path) == 2:
      print(datetime.now(), "RESULT FOR CUR_PATH = ", cur_path, "num_survivors = ", -best_path_eval[1], "prob_map = ", prob_map)
    if k > 0:
      return (prob_map[k], best_path_eval[1], prob_map)
    else:
      return (Fraction(0, 1), best_path_eval[1], prob_map)

def evaluate(n):
  return evaluate_dfs(n, {})

if __name__ == '__main__':
  print(datetime.now(), "Start")
  print(evaluate(int(sys.argv[1])))