Create an array of indices of values from a sequence of tuples when the values at each position in the tuples are from a different list

user_1609514 0
Viewed 177

This is hard to explain but I have a list of tuples where every tuple is length n and the possible values of the items in each tuple are contained in a set of lists (one list for each position in the tuple).

For the sake of a simple example, my list of tuples might be

sequence = [('b', 1), ('c', 2), ('c', 1), ('a', 3), ('c', 2), ('a', 3)]

The lists of possible values are:

state_lists = [['a', 'b', 'c'], [1, 2, 3, 4]]

I want an efficient way to create a list or array (integer array preferred) similar to sequence but with each value replaced by its index in the appropriate list from state_lists .

Here is one solution which utilizes np.searchsorted. Unfortunately, this function only supports 1-D arrays:

import numpy as np
sequence_as_array = np.array(sequence)
sequence_indexes = [
    np.searchsorted(states, sequence_as_array[:, i])
    for i, states in enumerate(state_lists)
]
lookup_array = np.vstack(sequence_indexes).T
print(lookup_array)

[[1 0]
 [2 1]
 [2 0]
 [0 2]
 [2 1]
 [0 2]]

Is there a simpler way to do this (e.g. without a for loop, conversion of sequence to an array and avoiding the vstack)?

It must be a general solution as the length of the tuples could be more than 2 and the data types could be float, int or string.

The reason for replacing the values is to create a lookup array for fast indexing and less memory than storing the original values. The length of sequence could be over 10,000.

tuples indexing numpy python
4 Answers

Use dicts for the reverse state to index maps?

state_dict = [{s: i for i,s in enumerate(state)} for state in state_lists]
lookup_array = np.array([[s[y] for s,y in zip(state_dict,x)] for x in sequence])
user_45659470

One efficient way to do this is to use the the built-in support for categorical sequences in pandas:

import pandas as pd
sequence = [('b', 1), ('c', 2), ('c', 1), ('a', 3), ('c', 2), ('a', 3)]
df = pd.DataFrame(sequence, dtype='category')

print(df[0].cat.categories.to_numpy())
# ['a', 'b', 'c']

print(df[0].cat.codes.to_numpy())
# [1 2 2 0 2 0]

Pandas has fairly optimized code-paths for this kind of thing that you can take advantage of for your use-case without having to re-implement them yourself.

user_29378310

Here's one solution I just discovered that avoids creating the array:

lookup_array = np.empty((len(sequence), len(state_lists)), dtype=int)
for i, item_seq in enumerate(zip(*sequence)):
    lookup_array[:, i] = np.searchsorted(state_lists[i], item_seq)
print(lookup_array)

[[1 0]
 [2 1]
 [2 0]
 [0 2]
 [2 1]
 [0 2]]
user_16095140

Taking a cue from @jakevdp's answer, we can also use the Pandas MultiIndex.from_tuples method to achieve the objective:

import pandas as pd

seq_as_multiindex = pd.MultiIndex.from_tuples(sequence)
lookup_array = np.array(seq_as_multiindex.codes).T
print(lookup_array)

[[1 0]
 [2 1]
 [2 0]
 [0 2]
 [2 1]
 [0 2]]

However, there is a caveat. This method does not use the state_lists to set up the categorical indexes so the values returned may not match. To avoid this, you can get the actual category labels used as follows:

print(seq_as_multiindex.levels)
[['a', 'b', 'c'], [1, 2, 3]]

(Note that this is slightly different from state_lists. There is no 4 in the second list).

user_16095140
Related