exchanges = {"A": ["AAPL", "BAC", "STW"],
           "B": ["BAC", "TSLA"]}
indices = {"A": ["AAPL", "BAC", "STW", "MSFT"],
             "B": ["BAC", "TSLA"],
             "C": ["AAPL", "BAC"],
             "D": ["AAPL", "BAC", "STW"]}

indexStocks = {} # dict that maps index to set of stocks in that index
for i in indices:
    indexStocks[i] = set()
    for stock in indices[i]:
        indexStocks[i].add(stock)

def computableIndices(stocks):
    # input: dictionary mapping stock name to 0 < number of occurrences, output set of computable indices given input exchanges
    # run time = O( (selected exchanges) * (stocks/exchange) + indices * (stocks/index) )
    
    out = set()
    for index in indices:
        status = True
        for stock in indexStocks[index]:
            if stock not in stocks:
                status = False
                break
        if status:
            out.add(index)
        
    return out


class Allocation:
    def __init__(self, exch_names):
        self.exch_names = exch_names
        
        self.stocks = {} # dict that maps stock name to number of occurrences in exchanges if stock was selected. Unselected stocks are not keys
        for name in self.exch_names:
            for stock in exchanges[name]:
                if stock in self.stocks:
                    self.stocks[stock] += 1
                else:
                    self.stocks[stock] = 1
        
        self.computableIndices = computableIndices(self.stocks) # set of computable indices

        self.missingStocks = {} # dict that maps index to number of stocks in that index that are not in the selected exchanges.
        for i in indices:
            self.missingStocks[i] = 0
            for stock in indexStocks[i]:
                if stock not in self.stocks: # unselected stock
                    self.missingStocks[i] += 1
        print(f"included stocks = {self.stocks}")
        print(f"missingStocks = {self.missingStocks}")
        
                
    def score(self):
        # must return number of indices computable with this allocation
        return len(self.computableIndices)
        
    def swap(self, ename_out, ename_in):
        # must change the allocation to remove ename_out and add ename_in.
        # we will run swap and then score in a loop 1e(4+) times so speed is a concern
        print(f"computable indices before swap = {self.computableIndices}. ")
        
        added = set(exchanges[ename_in]) # set of stocks added by the swap
        removed = set() # set of stocks that were removed after the swap
        for stock in exchanges[ename_out]:
            if stock not in added:
                if self.stocks[stock] == 1:
                    del self.stocks[stock]
                    removed.add(stock)
                else:
                    self.stocks[stock] -= 1
            else:
                added.remove(stock)
        print(f"added stocks = {added}")
        print(f"removed stocks = {removed}")

        for stock in added:
            self.stocks[stock] = 1
            for i in indices:
                if i not in self.computableIndices:
                    if stock in indexStocks[i]:
                        self.missingStocks[i] -= 1
                    if self.missingStocks[i] == 0:
                        self.computableIndices.add(i)

        for stock in removed:
            for i in self.computableIndices.copy():
                if stock in indexStocks[i]:
                    self.missingStocks[i] += 1
                    self.computableIndices.remove(i)
        print(f"computable indices after swap = {self.computableIndices}. ")

            
                            

a = Allocation(["B"])
a.score() #1

included stocks = {'BAC': 1, 'TSLA': 1}
missingStocks = {'A': 3, 'B': 0, 'C': 1, 'D': 2}

1

a.swap("B", "A")
a.score() #1

computable indices before swap = {'B'}. 
added stocks = {'STW', 'AAPL'}
removed stocks = {'TSLA'}
computable indices after swap = {'D', 'C'}. 

2

computableIndices(exchanges["A"])

{'C', 'D'}