Alternative Vote Counter

av.py

import csv
import re
import sys
import pprint

PRINT_VOTES = False
PRINT_OPTIONS = False

OPTION_PATTERN = r'.*\[(.*?)\].*'
VOTE_TEXTS = {
  'First': 1,
  'Second': 2,
  'Third': 3,
}


def GetVotes(file_path):
  votes = []
  with open(file_path) as f:
    reader = csv.reader(f)
    header = reader.next()
    option_cols = []
    option_names = []
    for i, col in enumerate(header):
      match = re.search(OPTION_PATTERN, col)
      if match:
        option_cols.append(i)
        option_names.append(match.group(1))
    for line in reader:
      vote = []
      for col in option_cols:
        vote_txt = line[col]
        for txt, num in VOTE_TEXTS.iteritems():
          if txt in vote_txt:
            vote_num = num
            break
        else:
          vote_num = int(vote_txt or '0')
        vote.append(vote_num)
      votes.append(vote)
  return option_names, votes

def TallyVotes(votes):
  options, votes = votes
  num_votes = len(votes)
  totals = [0] * len(options)
  for vote in votes:
    for i, num in enumerate(vote):
      if num == 1:
        totals[i] += 1
  totals_percent = [total / float(num_votes) for total in totals]
  PRINT_OPTIONS and pprint.pprint(zip(options, totals_percent))
  PRINT_VOTES and pprint.pprint(votes)
  if max(totals_percent) <= .5:
    min_percent = min([total for total in totals_percent if total > 0])
    min_cols = [i for i, v in enumerate(totals_percent) if v <= min_percent]
    max_vote = len(options) - len(min_cols)
    if max_vote <= 0:
      return options
    min_cols.reverse()
    new_votes = []
    new_options = options[:]
    for col in min_cols:
      del new_options[col]
    for vote in votes:
      new_vote = vote[:]
      for c, num in enumerate(vote):
        if c in min_cols:
          new_vote[c] = 0
        else:
          new_vote[c] = min(max_vote, new_vote[c])
      for col in min_cols:
        del new_vote[col]
      new_votes.append(new_vote)
    return TallyVotes((new_options, new_votes))
  else:
    return options

print TallyVotes(GetVotes(sys.argv[1]))