Python program to return rows that have element at a specified index

Last Updated : 18 May, 2023

Given two Matrices, the task is to write a Python program that can extract all the rows from both matrices which have similar elements at their Kth index, mapped at similar row positions.

Examples:

Input : test_list1 = [[1, 8, 3], [9, 2, 0], [6, 4, 4], [6, 4, 4]], 
        test_list2 = [[1, 9, 3], [8, 2, 3], [5, 4, 6], [5, 4, 6]], 
        K = 1
Output : [[9, 2, 0], [8, 2, 3], [6, 4, 4], [5, 4, 6], [6, 4, 4], [5, 4, 6]]
Explanation : All elements with similar elements at 1st index extracted.

Input : test_list1 = [[1, 8, 3], [9, 2, 0], [6, 4, 4], [6, 5, 4]], 
        test_list2 = [[1, 9, 3], [8, 2, 3], [5, 4, 6], [5, 4, 6]],
        K = 1
Output : [[9, 2, 0], [8, 2, 3], [6, 4, 4], [5, 4, 6]]
Explanation : All elements with similar elements at 1st index extracted.

Method 1: Using loop and enumerate()

In this, the list is iterated from the start row till the end, and each row’s Kth index is matched, if found, both rows are appended to the result.

Example:

Python3

# Initializing lists
test_list1 = [[1, 8, 3], [9, 2, 0], [6, 4, 4], [6, 4, 4]]
test_list2 = [[1, 9, 3], [8, 2, 3], [5, 4, 6], [5, 4, 6]]
 
# Printing original lists
print("The original list 1 is : " + str(test_list1))
print("The original list 2 is : " + str(test_list2))
 
# Initializing K
K = 1
 
# Empty list
res = []
 
for idx in range(len(test_list1)):
 
    # Comparinging lists
    if test_list1[idx][K] == test_list2[idx][K]:
        res.append(test_list1[idx])
        res.append(test_list2[idx])
 
 
# Printing result
print("K index matching rows : " + str(res))

Output

The original list 1 is : [[1, 8, 3], [9, 2, 0], [6, 4, 4], [6, 4, 4]]
The original list 2 is : [[1, 9, 3], [8, 2, 3], [5, 4, 6], [5, 4, 6]]
K index matching rows : [[9, 2, 0], [8, 2, 3], [6, 4, 4], [5, 4, 6], [6, 4, 4], [5, 4, 6]]

Time Complexity: O(n*m)
Auxiliary Space: O(k)

Method 2: Using list comprehension and zip()

In this, we perform the task of getting pairing using zip() and then compare the Kth element, append, and iterate using extend() and list comprehension.

Example:

Python3

# Initializing lists
test_list1 = [[1, 8, 3], [9, 2, 0], [6, 4, 4], [6, 4, 4]]
test_list2 = [[1, 9, 3], [8, 2, 3], [5, 4, 6], [5, 4, 6]]
 
# Printing original lists
print("The original list 1 is : " + str(test_list1))
print("The original list 2 is : " + str(test_list2))
 
# Initializing K
K = 1
 
# zip() combines elements together
res = []
[res.extend([t1, t2])
 for t1, t2 in zip(test_list1, test_list2) if t1[K] == t2[K]]
 
# Printing result
print("K index matching rows : " + str(res))

Output

The original list 1 is : [[1, 8, 3], [9, 2, 0], [6, 4, 4], [6, 4, 4]]
The original list 2 is : [[1, 9, 3], [8, 2, 3], [5, 4, 6], [5, 4, 6]]
K index matching rows : [[9, 2, 0], [8, 2, 3], [6, 4, 4], [5, 4, 6], [6, 4, 4], [5, 4, 6]]

Time Complexity: O(n*n) where n is the number of elements in the list “test_list”. The list comprehension and zip() are used to perform the task and it takes O(n*n) time.
Auxiliary Space: O(n), new list of size O(n) is created where n is the number of elements in the list.

Method 3: Using the NumPy module

Steps:

Import the numpy module.
Convert the input lists to numpy arrays.
Get the Kth column of each array using array slicing.
Compare the Kth columns of the arrays using the “==” operator to get a boolean array.
Use the boolean array to index the original arrays to get the matching rows.
Combine the matching rows from both arrays using NumPy’s “column_stack” function.
Convert the combined array back to a list of lists.

Python3

import numpy as np
 
# initializing lists
test_list1 = [[1, 8, 3], [9, 2, 0], [6, 4, 4], [6, 4, 4]]
test_list2 = [[1, 9, 3], [8, 2, 3], [5, 4, 6], [5, 4, 6]]
 
# convert lists to numpy arrays
arr1 = np.array(test_list1)
arr2 = np.array(test_list2)
 
# get Kth column of each array
k_col1 = arr1[:, 1]
k_col2 = arr2[:, 1]
 
# compare Kth columns of arrays to get boolean array
bool_arr = k_col1 == k_col2
 
# use boolean array to index original arrays to get matching rows
matching_arr1 = arr1[bool_arr]
matching_arr2 = arr2[bool_arr]
 
# combine matching rows from both arrays using vstack function
combined_arr = np.vstack((matching_arr1, matching_arr2))
 
# convert combined array back to list of lists
res = combined_arr.tolist()
 
# printing result
print("K index matching rows : " + str(res))

Output:

K index matching rows : [[9, 2, 0], [6, 4, 4], [6, 4, 4], [8, 2, 3], [5, 4, 6], [5, 4, 6]]

Time complexity: O(N), where N is the number of elements in the lists.
Auxiliary space: O(N) as well, since the arrays and the combined array both require N elements of memory.