def update_markov_blanket(MB, remove_nodes=[]):
    """
    Given a Markov Blanket dict and a (list of) node(s) that need to be
    updated in the blanket (i.e. they were recruited into a macro-node
    and should therefore not be searched), return a new Markov Blanket
    that has taken out the remove_nodes in question

    Parameters
    ----------
    MB (dict): Dictionary where the keys are the nodes in a graph, G,
               and the values are the nodes that constitute each node's
               Markov Blanket.
    remove_nodes (int or list): the nodes that have been recruited to
                    join a macro_node and that thus should not be
                    searched when trying to find other nodes to turn
                    into macro_nodes.

    Returns
    -------
    MB_new (list or dict): if remove_nodes was specified, this function
                       returns an updated Markov Blanket that lacks the
                       remove_nodes in the values of all the nodes.
                       Otherwise, it returns the original Markov Blanket.

    """

    if remove_nodes == []:
        return MB
    if type(remove_nodes) == int:
        remove_nodes = [remove_nodes]

    MB_new = {}
    for node_i, blanket_j in MB.items():
        MB_new[node_i] = [b for b in blanket_j if b not in remove_nodes]

    return MB_new