Open In App

Python – Retain all K elements Rows

Improve
Improve
Like Article
Like
Save
Share
Report

Sometimes, while working with Python lists, we can have a problem in which we need to retain rows which have only K as elements. This kind of application can occur in data domains which take Matrix as input. Let’s discuss certain ways in which this task can be performed.

Input : test_list = [[7, 6], [4, 4], [1, 2], [4]], K = 4 
Output : [[4, 4], [4]] 

Input : test_list = [[7, 6], [7, 4], [1, 2], [9]], K = 4 
Output : []

Method #1 : Using list comprehension + any() 

The combination of above functions provide a way in which this task can be performed. In this, we perform the task of filtering out rows which have any element other than K using any(). 

Python3




# Python3 code to demonstrate working of
# Retain all K elements Rows
# Using list comprehension + any()
 
# initializing list
test_list = [[2, 4, 6], [2, 2, 2], [2, 3], [2, 2]]
 
# printing original list
print("The original list : " + str(test_list))
 
# initializing K
K = 2
 
# Retain all K elements Rows
# Using list comprehension + any()
res = [ele for ele in test_list if not any(el != K for el in ele)]
 
# printing result
print("Matrix after filtering : " + str(res))


Output : 

The original list : [[2, 4, 6], [2, 2, 2], [2, 3], [2, 2]]
Matrix after filtering : [[2, 2, 2], [2, 2]]

time complexity : O(n*m)
Auxiliary space : O(k)

Method #2: Using list comprehension + all() The combination of above functions can be used to solve this problem. In this, we check if all the elements of row is equal to K using all(). 

Python3




# Python3 code to demonstrate working of
# Retain all K elements Rows
# Using list comprehension + all()
 
# initializing list
test_list = [[2, 4, 6], [2, 2, 2], [2, 3], [2, 2]]
 
# printing original list
print("The original list : " + str(test_list))
 
# initializing K
K = 2
 
# Retain all K elements Rows
# Using list comprehension + all()
res = [ele for ele in test_list if all(el == K for el in ele)]
 
# printing result
print("Matrix after filtering : " + str(res))


Output : 

The original list : [[2, 4, 6], [2, 2, 2], [2, 3], [2, 2]]
Matrix after filtering : [[2, 2, 2], [2, 2]]

time complexity : O(nm)
space complexity : O(nm)

 Method #3: Using filter() and lambda function:

Algorithm:

  1. Initialize an empty result list.
  2. Loop through each list in the input list.
  3. Check if all the elements in the current list are equal to K using the all() function.
  4. If all elements are equal to K, add the current list to the result list.
  5. Return the result list.
     

Python3




# Python3 code to demonstrate working of
# Retain all K elements Rows
# Using filter() and lambda function:
 
 
# initializing list
test_list = [[2, 4, 6], [2, 2, 2], [2, 3], [2, 2]]
 
# printing original list
print("The original list : " + str(test_list))
 
# initializing K
K = 2
 
# Retain all K elements Rows
# Using filter() and lambda function:
res = list(filter(lambda x: all(el == K for el in x), test_list))
 
# printing result
print("Matrix after filtering : " + str(res))
# this code contributed by tvsk


Output

The original list : [[2, 4, 6], [2, 2, 2], [2, 3], [2, 2]]
Matrix after filtering : [[2, 2, 2], [2, 2]]

Time complexity: O(n*m), where n is the number of lists in the input list and m is the maximum length of the lists.
Auxiliary space: O(k), where k is the number of lists that have all elements equal to K, since the result list will store at most k lists.

Method #4 : Using operator.countOf() method

Approach :

  1. Check whether the count of K  in row is equal to length of row using operator.countOf() method
  2. If yes then all elements are equal to K
  3. Append such rows to output list
  4. Display output list

Python3




# Python3 code to demonstrate working of
# Retain all K elements Rows
 
# initializing list
test_list = [[2, 4, 6], [2, 2, 2], [2, 3], [2, 2]]
 
# printing original list
print("The original list : " + str(test_list))
 
# initializing K
K = 2
 
# Retain all K elements Rows
res=[]
import operator
for i in test_list:
    if operator.countOf(i,K)==len(i):
        res.append(i)
 
# printing result
print("Matrix after filtering : " + str(res))


Output

The original list : [[2, 4, 6], [2, 2, 2], [2, 3], [2, 2]]
Matrix after filtering : [[2, 2, 2], [2, 2]]

Time Complexity : O(N)
Auxiliary Space : O(N)

Method #5: Using a for loop and a boolean flag variable to check if all elements in a row are equal to K.

This method iterates over each row of the matrix and checks if all elements in the row are equal to K using a boolean flag variable. If all elements in the row are equal to K, the row is added to the result list.

Python3




# Python3 code to demonstrate working of
# Retain all K elements Rows
 
# initializing list
test_list = [[2, 4, 6], [2, 2, 2], [2, 3], [2, 2]]
 
# printing original list
print("The original list : " + str(test_list))
 
# initializing K
K = 2
 
# Retain all K elements Rows
res=[]
for i in test_list:
    flag = True
    for j in i:
        if j != K:
            flag = False
            break
    if flag:
        res.append(i)
 
# printing result
print("Matrix after filtering : " + str(res))


Output

The original list : [[2, 4, 6], [2, 2, 2], [2, 3], [2, 2]]
Matrix after filtering : [[2, 2, 2], [2, 2]]

Time complexity: O(nm) where n is the number of rows and m is the number of columns in the matrix. 
Auxiliary space: O(k) where k is the number of rows that satisfy the condition.

Method 6:  Using list comprehension and zip()

  • Initialize the input matrix test_list and print it.
  • Initialize the value of K that we want to retain all rows for.
  • Create a list comprehension that filters the input matrix by checking if all elements in each row are equal to K.
  • Store the filtered rows in a new list called res.
  • Print the filtered matrix.

Python3




# initializing list
test_list = [[2, 4, 6], [2, 2, 2], [2, 3], [2, 2]]
 
# printing original list
print("The original list : " + str(test_list))
 
# initializing K
K = 2
 
# Retain all K elements Rows
# Using list comprehension and zip()
res = [row for row in test_list if all(x == K for x in row)]
 
# printing result
print("Matrix after filtering : " + str(res))


Output

The original list : [[2, 4, 6], [2, 2, 2], [2, 3], [2, 2]]
Matrix after filtering : [[2, 2, 2], [2, 2]]

Time complexity: O(n*m), where n is the number of rows and m is the maximum number of columns in the input list. This is because we need to iterate through each element of the array once to filter the rows.

Auxiliary space: O(1), as we are only creating a new list to store the filtered rows.

Method #7: Using numpy library

Import the numpy library.
Initialize a list named “test_list” containing sublists of different lengths.
Print the original list using the “print()” function.
Convert the list to a numpy array using the “numpy.array()” function.
Initialize a variable named “K” with the desired element.
Use the “numpy.where()” function to filter out rows from the numpy array where any element is not equal to “K”.
Store the filtered rows in a variable named “res”.
Convert the filtered rows to a list using the “tolist()” method of the numpy array.
Print the filtered list using the “print()” function.

Python3




# importing numpy library
import numpy as np
 
# initializing list
test_list = [[2, 4, 6], [2, 2, 2], [2, 3], [2, 2]]
 
# printing original list
print("The original list : " + str(test_list))
 
# converting list to numpy array
arr = np.array(test_list, dtype=object)
 
# initializing K
K = 2
 
# filtering rows where any element is not equal to K
res_arr = arr[np.where([(np.array(row) == K).all() for row in arr])]
 
# converting numpy array to list
res = res_arr.tolist()
 
# printing result
print("Matrix after filtering : " + str(res))


Output:

The original list : [[2, 4, 6], [2, 2, 2], [2, 3], [2, 2]]
Matrix after filtering : [[2, 2, 2], [2, 2]]

Time complexity: O(nm) where n is the number of rows and m is the number of columns in the matrix. 
Auxiliary space: O(k) where k is the number of rows that satisfy the condition.



Last Updated : 10 Apr, 2023
Like Article
Save Article
Previous
Next
Share your thoughts in the comments
Similar Reads