Convert a mapping of arbitrary pairs into a one-to-many map

Question

The accepted solution is much cleaner and outperforms the algorithm below for small maps, it doesn't do as well with larger ones: my code takes twice as long as the accepted code for maps of 10 pairs (30 vs. 15 ms), but mine is something like 30 times faster for the larger input described below (8 s vs. 4 m). However, @Reinderien's code never consumes more than around 120MB of memory. Mine is the same for 10-pair maps but it can take well over 1GB for large input (this is all in a Jupyter notebook on my intel i7 windows machine).

I've written an answer to this post with an improved algorithm that splits the difference - it handles the large test case in 55 seconds and consumes less than 600M of memory

The problem

I'm working in python. I have sets of pairs of integers that should be interpreted as bidirectional maps: values in one column should be mapped to values in the other column, but all paired values are considered equivalent. This means that the many-to-many map

src  tgt
  1    2
  2    3
  4    3
  4    1
  5    4

is equivalent to the one-to-many map

src  tgt
  5    1
  5    2
  5    3
  5    4

I want to take an arbitrary set of input pairs and generate a one-to-many map.

I'm happy with solutions using pandas or pure python, but solutions should only include the standard library and pandas.

Description of my solution

I've written a function that will convert an arbitrary two-column pandas dataframe of integers into a one-to-many map from the first column to the second column.

Below, "first instance" of a value means the earliest occurrence of a row with that value when scanning from the beginning of the dataframe in its current order (the lowest row index that will return the value when using df.iloc rather than df.loc).

There are about six steps in my solution:

1. Drop rows of the form [b, a] when the row [a, b] is also present
2. For values duplicated anywhere in the input regardless of column position,
   • 2.1 locate rows with the first second-column instance of any of the duplicate values
   • 2.2 locate rows with the first second-column instance of duplicate values whose first instance in the entire dataframe is in the second column
3. Map instances of the second-column values above to the first-column partners of their first instances in the second column
   • 3.1 for the set of rows 2.1, map all instances of the second-column values found anywhere in the first column to the first-column partners of the second-column values in set 2.1
   • 3.2 for the set of rows 2.2, map all duplicate instances of the second-column values in the second column to the first-column partner of the second-column values in set 2.2 and then swap the first and second columns in the rows with duplicate second-column values
4. Rearrange the input so that values duplicated anywhere in the input should appear first in the second column
5. If there are any rows whose second-column values also appear in the first column, swap the values in one of those rows (this is a kludge to avoid infinite loops)
6. If the map is not one-to-many, start over again

I expect there's an algorithm with fewer steps, but even if there isn't I think someone has to be able to do this more efficiently than I have. I have code to test performance at the end - I'm more interested in solutions that use less memory than solutions that run faster, but I'm interested in solutions that boost performance either way.

My code

Sorry for some obscenely long names; this has been a conceptual nightmare for me so I got verbose.

import pandas as pd

def make_bidirectional_map_one_to_many(df):

    input_columns = list(df.columns)
    df.columns = [0, 1]

    # Drop identities in a way that doesn't copy the whole dataframe.
    identities = df.loc[df[0] == df[1]]
    df = df.drop(identities.index).drop_duplicates()
    # Using the above method because the following line raises a copy
    # vs. view warning:
    # df = df.loc[df[0] != df[1], :]

    # Check if the input dataframe is already a one-to-many map
    if not df[1].duplicated().any() and not df[1].isin(df[0]).any():

        return df

    # Stack the values so we can find duplicate values regardless of the
    # column
    values = df.stack()
    duplicate_values = values.loc[values.duplicated(keep=False)]

    '''
    Drop rows of the form [b, a] when the row [a, b] is also present
    '''

    # Get values that are duplicated and present in both columns
    value_selection = duplicate_values.loc[
        duplicate_values.isin(df[0]) & duplicate_values.isin(df[1])
    ]
    drop_candidates = df.loc[df[1].isin(value_selection)]

    # Hash the dataframe to search for whole rows
    hashed_df = pd.util.hash_pandas_object(df, index=False)

    # Reverse the selected rows and hash them to search for the
    # reversed rows in the input dataframe
    reversed_candidates = drop_candidates[[1, 0]]
    hashed_candidates = pd.util.hash_pandas_object(
        reversed_candidates,
        index=False
    )

    # Drop the reversed rows if they're present in the dataframe
    drop_candidates = hashed_df.isin(hashed_candidates)
    if drop_candidates.any():

        # If the search above caught [a, b] then it also caught [b, a]
        # so we select one of those rows by stacking, sorting by index
        # and then by value, pivoting back to the original shape, and
        # then dropping duplicate rows.
        drop_candidates = df.loc[hashed_df[drop_candidates].index]
        drop_candidates = drop_candidates.stack().reset_index()
        drop_candidates = drop_candidates.sort_values(by=['level_0', 0])
        drop_candidates.loc[:, 'level_1'] = [0, 1]*int(len(drop_candidates)/2)
        drop_candidates = drop_candidates.pivot(
            index='level_0',
            columns='level_1',
            values=0
        )

        drop_rows = drop_candidates.drop_duplicates().index

        df = df.drop(drop_rows)

        # If the dataframe is now a one-to-many map, we're done
        if not df[1].duplicated().any():
            if not df[1].isin(df[0]).any():
                return df

        # Stack the modified dataframe and get duplicate values
        values = df.stack()
        duplicate_values = values.loc[values.duplicated(keep=False)]

    '''
    Locate two different types of duplicate values in the second column
    '''

    # Get first instance in the second column of values duplicated
    # anywhere in the dataframe
    first_instance_of_any_dplct_in_c1 = (
        duplicate_values.loc[(slice(None), 1)].drop_duplicates()
    )

    # Get values from the second column which are duplicated anywhere
    # in the dataframe but which occurred first in the second column
    frst_instnc_of_dplct_found_in_c1_frst = duplicate_values.drop_duplicates()
    if 1 in frst_instnc_of_dplct_found_in_c1_frst.index.get_level_values(1):
        frst_instnc_of_dplct_found_in_c1_frst = (
            frst_instnc_of_dplct_found_in_c1_frst.loc[(slice(None), 1)]
        )

    else:
        frst_instnc_of_dplct_found_in_c1_frst = pd.Series([], dtype=int)

    # If values duplicated anywhere in the dataframe exist in the
    # second column, map the duplicates to the first-column
    # partner of the first instance of the value in the second column
    if df[1].duplicated().any():

        '''
        Make maps between the different kinds of first second-column
        instances and their first-column partners
        '''

        # Map from the first instance in the second column of values
        # duplicated anywhere in the dataframe to their first-column
        # partners
        df_selection = df.loc[first_instance_of_any_dplct_in_c1.index]
        map_to_c0_from_frst_instnc_of_any_dplct_in_c1 = pd.Series(
            df_selection[0].array,
            index=df_selection[1].array
        )

        # Map from duplicate values that occurred first in the second
        # column to the first-column partners of their first instance
        map_to_c0_from_frst_instnc_of_dplct_found_in_c1_frst = df.loc[
            frst_instnc_of_dplct_found_in_c1_frst.index,
            :
        ]
        map_to_c0_from_frst_instnc_of_dplct_found_in_c1_frst = pd.Series(
            map_to_c0_from_frst_instnc_of_dplct_found_in_c1_frst[0].array,
            index=map_to_c0_from_frst_instnc_of_dplct_found_in_c1_frst[1].array
        )

        '''
        Apply the maps
        '''

        # Get rows with values in the first column that appeared
        # first in the second column
        c0_value_was_found_in_c1_first = df.loc[
            df[0].isin(first_instance_of_any_dplct_in_c1),
            :
        ]

        # Map those values to the first-column partners of their first
        # instances in the second column
        mapped_c0 = c0_value_was_found_in_c1_first[0].map(
            map_to_c0_from_frst_instnc_of_any_dplct_in_c1
        )
        df.loc[c0_value_was_found_in_c1_first.index, 0] = mapped_c0.array

        # Get all duplicate values in the second column which appeared first
        # in the second column
        c1_value_was_found_in_c1_first = df.loc[
            df[1].isin(frst_instnc_of_dplct_found_in_c1_frst),
            :
        ]

        # Drop the first instances to isolate only the duplicates
        c1_value_was_found_in_c1_first = c1_value_was_found_in_c1_first.drop(
            frst_instnc_of_dplct_found_in_c1_frst.index
        )

        # If there are duplicate values in the second column which
        # appeared first in the second column, map those values to the
        # first-column partners of their first instances and then swap
        # values in those rows so the first instance's first-column
        # partner is in the first column

        if not c1_value_was_found_in_c1_first.empty:
            mapped_1 = c1_value_was_found_in_c1_first[1].map(
                map_to_c0_from_frst_instnc_of_dplct_found_in_c1_frst
            )
            df.loc[c1_value_was_found_in_c1_first.index, :] = list(zip(
                mapped_1.array,
                df.loc[c1_value_was_found_in_c1_first.index, 0]
            ))
            df = df.drop_duplicates()

    # Rearrange the dataframe so that values which appear in both
    # columns appear first in the second column

    c1_in_c0 = df[1].isin(df[0])
    df = pd.concat([df.loc[c1_in_c0], df.loc[~c1_in_c0]])
    
    c0_in_c1 = df[0].isin(df[1])
    df = pd.concat([df.loc[~c0_in_c1], df.loc[c0_in_c1]])

    # In some cases the above algorithm will loop indefinitely if there
    # are values that occur in both columns but there are no repeated
    # values in the second column. Avoid this finding rows with values
    # in the first column that are present in the second column and, if
    # any exist, swapping values in the first row found.
    c0_in_c1 = df[0].isin(df[1])

    if c0_in_c1.any():
        to_swap = df.loc[c0_in_c1].index[0]
        df.loc[to_swap, :] = df.loc[to_swap, [1, 0]].to_numpy()

    df.columns = input_columns

    # Remap the modified dataframe
    return make_bidirectional_map_one_to_many(df)

Code to test solutions

The above code gives this result

a = pd.DataFrame({0:[1,2,4,4,5], 1:[2,3,3,1,4]})
make_bidirectional_map_one_to_many(a)

    0   1
2   5   2
1   5   3
0   5   4
3   5   1

But note that since the map is bidirectional, it doesn't matter which integer is in column 0 and the ordering in the output doesn't matter.

I can run the following code from a VS Code jupyter notebook (Windows) in about 32 seconds on an intel i7, and the python process in the task manager never consumes more than about 100M of memory. With number_of_runs = 10, pop_size = 1e4, and sample_size = 1e6, it takes about 81 seconds and memory consumption can spike to 1.3G but probably averages about 800M.

import time
import pandas as pd

number_of_runs = 1e3

def run_it():
    pop_size = 1e1
    sample_size = 1e1

    t = str(time.time()).split('.')
    s1 = int(t[0])
    s2 = int(t[1])
    pop = pd.Series(pd.RangeIndex(stop=pop_size))
    sampled = pd.DataFrame({
        'a': pop.sample(int(sample_size), replace=True, random_state=s1).array,
        'b': pop.sample(int(sample_size), replace=True, random_state=s2).array
    })

    try:
        make_bidirectional_map_one_to_many(sampled)

    except:
        print(sampled)
        raise RuntimeError

for i in range(1, int(number_of_runs) + 1):
    time.sleep(.005)
    run_it()

Answer 1

Somewhat guessing what you're really trying to do:

Forget Pandas for the moment. Consider the left column to be the first vertices of a graph edge, and the second column to contain the second vertex for each of those edges. Your graph has properties:

• Undirected
• Finite
• Not necessarily acyclic, connected or complete

You want to apply a graph transformation to produce a new graph where:

• You identify components of your original graph
• Within each component, you
  • Choose (arbitrarily?) a vertex to become universal
  • Remove all other edges.

This should give you enough Google fodder. I cannot promise that the following suggested algorithm is the most efficient, but it approaches something sane. It drops out of Pandas for simplicity's sake but certain parts may be further vectorisable.

import pandas as pd


def make_bidirectional_map_one_to_many(df: pd.DataFrame) -> pd.DataFrame:
    # Dictionary of every known key to a set
    # representing the component we think it's in.
    components: dict[int, set[int]] = {}

    for idx, (src_k, tgt_k) in df.iterrows():
        # Start off with either existing component sets, or make singleton sets.
        src_set = components.setdefault(src_k, {src_k})
        tgt_set = components.setdefault(tgt_k, {tgt_k})

        # Overwrite all of the source set entries with a reference to the target set
        components.update({
            k: tgt_set for k in tuple(src_set)
        })

        # Union the source set into the target set
        # (all references to this set will update)
        tgt_set |= src_set

    # All vertices that we've added to the universals list so far
    covered = set()
    # List of source-target tuples
    universals = []

    for k, component in components.items():
        if k not in covered:
            covered |= component
            for c in component:
                if k != c:
                    universals.append((k, c))

    return pd.DataFrame.from_records(universals, columns=('src', 'tgt'))

Answer 2

I've cut the memory consumption by about half and the run time for the large test case described in the OP is now about 55 seconds (5.5 seconds per million-pair dataframe). Peak memory consumption in the windows task manager on an intel i7 machine running the code in a VS code jupyter notebook is less than 600M for that test.

Performance improvement achieved in three ways:

1. Replacing a bunch of variables with dummy. This obviously makes the code less readable but saves a couple of hundred megabytes of memory
2. Using inplace=True when convenient
3. Eliminating the part of the algorithm that drops redundant edges. There was a sort operation in there that's pretty costly. This might create some tradeoffs for dataframes that have a lot of duplicates, but I don't know.

Note also that upgrading to pandas 1.5 might have significance here. For instance statements like df.drop_duplicates() used to be more resource-intensive than they appear to be now, but I don't know which version might have made that happen (I've been avoiding some things for quite a while and just found some improvements playing around with this code today)

import pandas as pd

def make_bidirectional_map_one_to_many(df):

    input_columns = list(df.columns)
    df.columns = [0, 1]

    # Remove identities and duplicate rows
    df = df.loc[df[0] != df[1], :].drop_duplicates()

    # Check if the input dataframe is already a one-to-many map
    if not df[1].duplicated().any() and not df[1].isin(df[0]).any():
        df.columns = input_columns
        return df

    # Stack the values so we can find duplicate values regardless of the
    # column
    dummy = df.stack()
    dummy = dummy.loc[dummy.duplicated(keep=False)]

    ##############################################################
    # Locate two different types of duplicate values in the second
    # column
    ##############################################################

    # Get first instance in the second column of values duplicated
    # anywhere in the dataframe. "dummy" refers to duplicate values in
    # the stacked dataframe
    first_instance_of_any_dplct_in_c1 = (
        dummy.loc[(slice(None), 1)].drop_duplicates()
    )

    # Get values from the second column which are duplicated anywhere
    # in the dataframe but which occurred first in the second column.
    # "dummy" refers to duplicate values in the stacked dataframe
    frst_instnc_of_dplct_found_in_c1_frst = dummy.drop_duplicates()
    if 1 in frst_instnc_of_dplct_found_in_c1_frst.index.get_level_values(1):
        frst_instnc_of_dplct_found_in_c1_frst = (
            frst_instnc_of_dplct_found_in_c1_frst.loc[(slice(None), 1)]
        )

    else:
        frst_instnc_of_dplct_found_in_c1_frst = pd.Series([], dtype=int)

    # If values duplicated anywhere in the dataframe exist in the
    # second column, map the duplicates to the first-column
    # partner of the first instance of the value in the second column
    if df[1].duplicated().any():

        ##############################################################
        # Make maps between the different kinds of first second-column
        # instances and their first-column partners
        ##############################################################

        # Map from the first instance in the second column of values
        # duplicated anywhere in the dataframe to their first-column
        # partners.
        # "dummy" refers to duplicate values in the stacked dataframe
        dummy = df.loc[first_instance_of_any_dplct_in_c1.index]
        map_to_c0_from_frst_instnc_of_any_dplct_in_c1 = pd.Series(
            dummy[0].array,
            index=dummy[1].array
        )

        # Map from duplicate values that occurred first in the second
        # column to the first-column partners of their first instance
        map_to_c0_from_frst_instnc_of_dplct_found_in_c1_frst = df.loc[
            frst_instnc_of_dplct_found_in_c1_frst.index,
            :
        ]
        map_to_c0_from_frst_instnc_of_dplct_found_in_c1_frst = pd.Series(
            map_to_c0_from_frst_instnc_of_dplct_found_in_c1_frst[0].array,
            index=map_to_c0_from_frst_instnc_of_dplct_found_in_c1_frst[1].array
        )

        ################
        # Apply the maps
        ################

        # Get rows with values in the first column that appeared
        # first in the second column
        c0_value_was_found_in_c1_first = df.loc[
            df[0].isin(first_instance_of_any_dplct_in_c1),
            :
        ]

        # Map those values to the first-column partners of their first
        # instances in the second column
        mapped = c0_value_was_found_in_c1_first[0].map(
            map_to_c0_from_frst_instnc_of_any_dplct_in_c1
        )
        df.loc[c0_value_was_found_in_c1_first.index, 0] = mapped.array

        # Get all duplicate values in the second column which appeared first
        # in the second column
        c1_value_was_found_in_c1_first = df.loc[
            df[1].isin(frst_instnc_of_dplct_found_in_c1_frst),
            :
        ]

        # Drop the first instances to isolate only the duplicates
        c1_value_was_found_in_c1_first = c1_value_was_found_in_c1_first.drop(
            frst_instnc_of_dplct_found_in_c1_frst.index
        )

        # If there are duplicate values in the second column which
        # appeared first in the second column, map those values to the
        # first-column partners of their first instances and then swap
        # values in those rows so the first instance's first-column
        # partner is in the first column

        if not c1_value_was_found_in_c1_first.empty:
            mapped = c1_value_was_found_in_c1_first[1].map(
                map_to_c0_from_frst_instnc_of_dplct_found_in_c1_frst
            )
            df.loc[c1_value_was_found_in_c1_first.index, :] = list(zip(
                mapped.array,
                df.loc[c1_value_was_found_in_c1_first.index, 0]
            ))
            df.drop_duplicates(inplace=True)
        
        # Remove any identities or duplicate rows created above
        df = df.loc[df[0] != df[1], :].drop_duplicates()

    # Rearrange the dataframe so that values which appear in both
    # columns appear first in the second column

    dummy = df[1].isin(df[0])
    df = pd.concat([df.loc[dummy], df.loc[~dummy]])

    dummy = df[0].isin(df[1])
    df = pd.concat([df.loc[~dummy], df.loc[dummy]])

    # In some cases the above algorithm will loop indefinitely if there
    # are values that occur in both columns but there are no repeated
    # values in the second column. Avoid this finding rows with values
    # in the first column that are present in the second column and, if
    # any exist, swapping values in the first row found.
    dummy = df[0].isin(df[1])

    if dummy.any():
        to_swap = df.loc[dummy].index[0]
        df.loc[to_swap, :] = df.loc[to_swap, [1, 0]].to_numpy()

    df.columns = input_columns

    # Remap the modified dataframe
    return make_bidirectional_map_one_to_many(df)

df = pd.DataFrame({'tail': [1, 4, 2, 4, 6, 5, 7, 8], 'head': [2, 1, 3, 3, 8, 4, 6, 7]})
make_bidirectional_map_one_to_many(df)
```