Python – Extract ith column values from jth column values

Sometimes, while working with Python Matrix, we can have a problem in which we need to extract ith column values from comparing values from jth column. This kind of problem can occur in domains such as school programming or web development. Let’s discuss certain ways in which this task can be performed.

Input : test_list = [[4, 5, 6], [2, 5, 7], [9, 8, 2], [10, 2, 6]], 
        search_list = [4, 9], search_idx = 0, ext_idx = 2 
Output : [6, 2] 

Input : test_list = [[4, 5, 6], [2, 5, 7], [9, 8, 2], [10, 2, 6]], 
        search_list = [2, 6], search_idx = 2, ext_idx = 0 
Output : [4, 9, 10]

Method #1: Using loop

This is a brute way to solve this problem. In this, we loop through each row and compare the jth column with list elements, if present, extract the ith element. 

Python3

# Python3 code to demonstrate working of
# Extract ith column values from jth column values
# Using loop

# initializing list
test_list = [[4, 5, 6], [2, 5, 7], [9, 8, 2], [10, 2, 6]]

# printing original list
print("The original list is : " + str(test_list))

# initializing list
search_list = [5, 2]

# initializing search index
search_idx = 1

# initializing extract index
ext_idx = 2

# Extract ith column values from jth column values
# Using loop
res = []
for sub in test_list:
    if sub[search_idx] in search_list:
        res.append(sub[ext_idx])

# printing result
print("The extracted elements : " + str(res))

Output : 

The original list is : [[4, 5, 6], [2, 5, 7], [9, 8, 2], [10, 2, 6]]
The extracted elements : [6, 7, 6]

Time complexity: O(n), where n is the number of sub-lists in the test_list. This is because the loop runs once for each sub-list in the test_list.
Auxiliary Space: O(m), where m is the number of sub-lists in test_list that meet the condition. This is because the list res is created to store the extracted elements and its length depends on the number of sub-lists that meet the condition.

Method #2: Using set() + list comprehension The combination of the above functions can be used to solve this problem. In this, we perform the task of extracting elements in a similar way as the above method, just using shorthand. Also, for initial data search reduction, the search list is converted to set(), to remove duplicates. 

Python3

# Python3 code to demonstrate working of
# Extract ith column values from jth column values
# Using set() + list comprehension

# initializing list
test_list = [[4, 5, 6], [2, 5, 7], [9, 8, 2], [10, 2, 6]]

# printing original list
print("The original list is : " + str(test_list))

# initializing list
search_list = [5, 2]

# initializing search index
search_idx = 1

# initializing extract index
ext_idx = 2

# Extract ith column values from jth column values
# Using set() + list comprehension
temp = set(search_list)
res = [sub[ext_idx] for sub in test_list if sub[search_idx] in search_list]

# printing result
print("The extracted elements : " + str(res))

Output : 

The original list is : [[4, 5, 6], [2, 5, 7], [9, 8, 2], [10, 2, 6]]
The extracted elements : [6, 7, 6]

Time complexity: O(n), where n is the number of elements in the test_list.
Auxiliary space: O(k), where k is the number of elements in the search_list.

Method 3-using the filter() function along with a lambda function

This implementation first uses filter() along with a lambda function to filter out the rows that contain the search values in the given column index. Then, it uses a list comprehension to extract the required values from the filtered list.

Python3

# Python3 code to demonstrate working of
# Extract ith column values from jth column values
# Using filter() and lambda

# initializing list
test_list = [[4, 5, 6], [2, 5, 7], [9, 8, 2], [10, 2, 6]]

# printing original list
print("The original list is : " + str(test_list))

# initializing list
search_list = [5, 2]

# initializing search index
search_idx = 1

# initializing extract index
ext_idx = 2

# Extract ith column values from jth column values
# Using filter() and lambda
res = list(filter(lambda x: x[search_idx] in search_list, test_list))
res = [sub[ext_idx] for sub in res]

# printing result
print("The extracted elements : " + str(res))

Output

The original list is : [[4, 5, 6], [2, 5, 7], [9, 8, 2], [10, 2, 6]]
The extracted elements : [6, 7, 6]

Time complexity: O(n), where n is the number of elements in the input list. 
Auxiliary space: O(k), where k is the number of elements in the resulting list. 

Method 4: Using list comprehension without set()

Python3

test_list = [[4, 5, 6], [2, 5, 7], [9, 8, 2], [10, 2, 6]]
search_list = [5, 2]
search_idx = 1
ext_idx = 2

res = [sub[ext_idx] for sub in test_list if sub[search_idx] in search_list]
print("The extracted elements : " + str(res))

Output

The extracted elements : [6, 7, 6]

Time complexity: O(n), where n is the number of sub-lists in test_list.
Auxiliary space: O(k), where k is the number of sub-lists in test_list.

Method 5: Using numpy

You can use the NumPy library to extract the ith column values from jth column values in a more efficient way.

Approach:

1. Import the NumPy library.
2. Convert the given list into a NumPy array using the numpy.array() function.
3. Use NumPy indexing to extract the ith and jth columns.
4. Use the NumPy where() function to get the indices where the jth column values match the values in the search_list.
5. Use the extracted indices to get the ith column values.
6. Convert the extracted values back to a Python list using the tolist() function.
7. Print the extracted elements.

Python3

# Python3 code to demonstrate working of
# Extract ith column values from jth column values
# Using numpy

# import numpy library
import numpy as np

# initializing list
test_list = [[4, 5, 6], [2, 5, 7], [9, 8, 2], [10, 2, 6]]

# printing original list
print("The original list is : " + str(test_list))

# initializing list
search_list = [5, 2]

# initializing search index
search_idx = 1

# initializing extract index
ext_idx = 2

# convert list to numpy array
test_arr = np.array(test_list)

# extract ith and jth columns using numpy indexing
ij_cols = test_arr[:, [search_idx, ext_idx]]

# get the indices where jth column values match the search_list
indices = np.where(np.isin(ij_cols[:, 0], search_list))

# extract the ith column values using the indices
res = ij_cols[indices, 1].tolist()

# flatten the result list
res = [val for sublist in res for val in sublist]

# printing result
print("The extracted elements : " + str(res))

OUTPUT:

The original list is : [[4, 5, 6], [2, 5, 7], [9, 8, 2], [10, 2, 6]]
The extracted elements : [6, 7, 6]

Time complexity: O(N), where N is the number of elements in the given list.
Auxiliary space: O(N), where N is the number of elements in the given list.

Method 6: Using pandas

We can also use the pandas library to extract the ith column values from the jth column values. Here’s the step-by-step approach:

1. Import the pandas library.
2. Convert the test_list list to a pandas DataFrame using the pd.DataFrame() function.
3. Use the loc[] function to select the jth column values where the values are in search_list.
4. Use the iloc[] function to select the ith column values from the jth column values selected in step 3.
5. Store the result in res.
6. Print the result.

Python3

# Import the pandas library
import pandas as pd

# initializing list
test_list = [[4, 5, 6], [2, 5, 7], [9, 8, 2], [10, 2, 6]]

# printing original list
print("The original list is : " + str(test_list))

# initializing list
search_list = [5, 2]

# initializing search index
search_idx = 1

# initializing extract index
ext_idx = 2

# Convert the test_list to a pandas DataFrame
df = pd.DataFrame(test_list)

# Use loc[] to select the jth column values where the values are in search_list
jth_col = df.loc[df[search_idx].isin(search_list), search_idx]

# Use iloc[] to select the ith column values from df
res = df.iloc[jth_col.index, ext_idx].tolist()

# printing result
print("The extracted elements : " + str(res))

OUTPUT:
The original list is : [[4, 5, 6], [2, 5, 7], [9, 8, 2], [10, 2, 6]]
The extracted elements : [6, 7, 6]

Time complexity: O(n), where n is the number of elements in the test_list list.
Auxiliary space: O(n), where n is the number of elements in the test_list list, due to the conversion of the list to a pandas DataFrame.