UFC fights metadata

test.py

import csv

fighterToWeight = {}
weightClasses = ["Atomweight", "Strawweight", "Flyweight",
                 "Bantamweight", "Featherweight", "Lightweight",
                 "Welterweight", "Middleweight", "Light Heavyweight",
                 "Heavyweight", "Super Heavyweight"]

# Set up info about weight classes
with open('fighters.csv') as table:
    reader = csv.DictReader(table)

    for row in reader:
        fighterToWeight[row["name"]] = row["class"]


# Maps fighter name -> number of fights
fighters = {}

thrownOutData = 0
fightsAcrossWeightClasses = 0

# Read in each fight. Keep track of # of fights
# that each fighter has been in
with open('fights.csv') as table:
    reader = csv.DictReader(table)

    for row in reader:
        name1 = row['f1name']
        name2 = row['f1name']

        # We don't have full info on one of the participating fighters
        # so exclude this fight from the data we consider
        if name1 not in fighterToWeight or \
        name2 not in fighterToWeight or \
        fighterToWeight[name1] not in weightClasses or \
        fighterToWeight[name2] not in weightClasses:
            thrownOutData += 1
            continue
        
        if name1 in fighters:
            fighters[name1] += 1
        else:
            fighters[name1] = 0
            
        if name2 in fighters:
            fighters[name2] += 1
        else:
            fighters[name2] = 0

        if fighterToWeight[name1] != fighterToWeight[name2]:
            fightsAcrossWeightClasses += 1

with open('out.csv', 'w') as output:
    # Write first row (header)
    output.write("Min # of Fights,")

    for class_ in weightClasses:
        output.write(class_ + ",")
        
    output.write("Total\n")

    # Write a row for each minimum # of fights
    for threshold in range(1, 39):
        fightersPerWeightClass = {}
        total = 0
        
        for class_ in weightClasses:
            fightersPerWeightClass[class_] = 0
            
        for name, fights in fighters.items():
            if(fights < threshold):
                continue

            fightersPerWeightClass[fighterToWeight[name]] += 1
            total += 1

        output.write(str(threshold) + ",")
        for class_ in weightClasses:
            output.write(str(fightersPerWeightClass[class_]) + ",")
        output.write(str(total) + "\n")
                
print(str(fightsAcrossWeightClasses) + " fights across weight classes.")
print("Threw out " + str(thrownOutData) + " fights due to malformed data.")
print("Printing table to out.csv")
print("")
print("*** Complete! ***")