return_affiliation_df.py

import pandas as pd
import numpy as np
from itertools import combinations, permutations
from collections import defaultdict
import matplotlib.pyplot as plt
from collections import Counter
import seaborn as sns

def return_affiliation_df(dd):
    """
        Returns an affiliation index data frame per voting round
        input: dd - data frame of voting history
        output: affiliation_dict_list - data frame of affiliation index per voting round
    """
    flatten = lambda l: [item for sublist in l for item in sublist]
    contestants = dd.index
    pairs = list(combinations(contestants,2))
    affiliation_dict_list = []
    
    # iterate over voting rounds
    for col_index in dd.columns:
        # get the clusters who voted for particular people
        voting_dict = dd.groupby(col_index).groups
        # if not a name, remove from the voting dict
        for i in ['','—']:
            if i in voting_dict.keys():
                voting_dict.pop(i)
        # everyone who voted in tribal council
        in_tribal = set(flatten(voting_dict.values()))
        # get all the people who received votes at tribal
        candidates = voting_dict.keys()
        candidates = [name for name in candidates if name in contestants]
                
        # get the vote counts of the people who received votes
        vote_counts = []
        for voted in candidates:
            vote_counts.append(len(voting_dict[voted]))
 
        affiliation_dict = defaultdict(int)
        # initialize all pairs of people in tribal to have affiliation value of 0
        for i in in_tribal:
            for j in in_tribal:
                if (i,j) in pairs:
                    affiliation_dict[(i,j)] = 0
        # in the case of tiebreakers, or when only 2 people are at risk, set affiliation value
        # of a pair of contestants to 1 if they voted together
        if 'None' in voting_dict.keys() or len(candidates) < 3:
            for voted in candidates:
                for i in voting_dict[voted]:
                    for j in voting_dict[voted]:
                        if (i,j) in pairs:
                            affiliation_dict[(i,j)] = 1 
            affiliation_dict_list.append(affiliation_dict)
            continue
        # in the case of split votes:
        else:
            # there are usually 3 names thrown out during split votes:
            # the 2 targets of the majority affiliation, and the target of the
            # people in the minority.

            # get who the contestants at risk voted for
            candidates_voted_for = dict()
            for voted in candidates:
                for other in candidates:
                    if voted in voting_dict[other]:
                        candidates_voted_for[voted]=other
            #print(col_index,candidates_voted_for)
                        
            # cross contains all possible 3 pairs of the contestants at risk
            cross = list(combinations(candidates,3))
            # cluster_i represents the contestants who voted for contestant i
            for (voted1, voted2, voted3) in cross:
                cluster1 = voting_dict[voted1]
                cluster2 = voting_dict[voted2]
                cluster3 = voting_dict[voted3]

                # set affiliation value of a pair of contestants to 1 if the they
                # belong to the same cluster
                for i in cluster1:
                    for j in cluster1:
                        if (i,j) in pairs:
                            affiliation_dict[(i,j)] = 1

                for i in cluster2:
                    for j in cluster2:
                        if (i,j) in pairs:
                            affiliation_dict[(i,j)] = 1

                for i in cluster3:
                    for j in cluster3:
                        if (i,j) in pairs:
                            affiliation_dict[(i,j)] = 1

                # target_i is the vote of the ith contestant at risk
                target1 = candidates_voted_for[voted1]
                target2 = candidates_voted_for[voted2]
                target3 = candidates_voted_for[voted3]

                voteds = [voted1,voted2,voted3]
                clusters = [cluster1, cluster2, cluster3]
                targets = [target1, target2, target3]
                index_pairs = combinations([0,1,2],2)
                for (ix1, ix2) in index_pairs:
                    for cluster in clusters:
                        # check if more than 1 person votes for contestant ix1 and more than
                        # 1 person votes for contestant ix2, then they are candidates for the
                        # split vote.
                        if (len(clusters[ix1]) > 1) and (len(clusters[ix2]) > 1):
                            # check if the targets of contenstant ix1 and ix2 are in the same cluster
                            # if yes, then the people who voted for ix1 and ix2 did a split
                            # vote and are in an alliance
                            if (targets[ix1] in cluster) and (targets[ix2] in cluster):
                                everyone = flatten([clusters[ix1],clusters[ix2]])
                                for i in everyone:
                                    for j in everyone:
                                        if (i,j) in pairs:
                                            # assign affiliation = 1 to all who voted for 
                                            # ix1 and ix2
                                            affiliation_dict[(i,j)] = 1
                                # since ix1 and ix2 voted together, they are in a
                                # minority affiliation
                                if (voted[ix1], voted[ix2]) in pairs:
                                    affiliation_dict[(voted[ix1], voted[ix2])] = 1
                                elif (voted[ix2], voted[ix1]) in pairs:
                                    affiliation_dict[(voted[ix2], voted[ix1])] = 1
            affiliation_dict_list.append(affiliation_dict)
    return pd.DataFrame(affiliation_dict_list).T