Python – Next N elements from K value

Given a List, get next N values from occurrence of K value in List.

Input : test_list = [3, 4, 6, 7, 8, 4, 7, 2, 1, 8, 4, 2, 3, 9], N = 1, K = 4 
Output : [6, 7, 2] 
Explanation : All successive elements to 4 are extracted as N = 1.

Input : test_list = [3, 4, 6, 7, 8, 4, 7, 2, 1, 8, 4, 2, 3, 9], N = 2, K = 7 
Output : [8, 4, 2, 1] 
Explanation : Two elements after each occurrence of 7 are extracted. 

Method #1 : Using list comprehension + slicing 

In this we extract all the indices using list comprehension, and then loop is employed to append the elements and join in as single string from the desired K value occurrences in List.

The original list is : [3, 4, 6, 7, 8, 4, 7, 2, 1, 8, 4, 2, 3, 9]
Constructed Result List : [6, 7, 7, 2, 2, 3]

Time Complexity: O(n*n) where n is the number of elements in the list “test_list”.  list comprehension + slicing performs n*n number of operations.
Auxiliary Space: O(n), extra space is required where n is the number of elements in the list

Method #2 : Using filter() + lambda + loop

This is similar to above method, difference being filtering of indices is performed using filter() and lambda functions.

The original list is : [3, 4, 6, 7, 8, 4, 7, 2, 1, 8, 4, 2, 3, 9]
Constructed Result List : [6, 7, 7, 2, 2, 3]

Method #3 : Using enumerate

Approach

we can use the enumerate() function to iterate over the elements of the input list and find the index of the specified value. Then, we can use slicing to get the next N elements after the specified value.

Algorithm

1. Initialize an empty list called result.
2. Iterate over the input list using enumerate() function to get the index and value of each element.
3. If the value at the current index is equal to the specified value K, use slicing to get the next N elements after the specified value, and extend the result list with the sliced list.
4. Return the result list.

[6, 7, 7, 2, 2, 3]

Time complexity: O(N*M), where N is the number of occurrences of the specified value K in the input list and M is the size of the sliced list. In the worst case, all the elements in the input list are equal to K, and we need to slice the entire list, which takes O(M) time.

Auxiliary Space: O(M), where M is the size of the sliced list. In the worst case, if we need to slice the entire list, the space complexity will be O(NM). However, in practice, the size of the sliced list is usually much smaller than the size of the input list, so the space complexity is usually much smaller than O(NM).

Method 4: Using a generator function

Initialize a list called test_list with some integer values.
Print the original list test_list using the print() function.
Initialize an integer variable K with a value of 4.
Initialize an integer variable N with a value of 2.
Define a generator function called get_next_n_elements_after_k that takes three arguments: a list lst, an integer k, and an integer n.
The generator function uses a for loop to iterate through each element of the list lst along with its index using the enumerate() function.
If the current element in the loop is equal to the integer k, the generator function yields the next n elements after k using slicing with the yield keyword.
Outside of the generator function, initialize an empty list called res.
Use a for loop to iterate through each group of n elements after k using the get_next_n_elements_after_k generator function with arguments test_list, K, and N.
For each group of n elements, extend the res list with the group of elements using the extend() method.
Print the constructed result list res using the print() function.

The original list is : [3, 4, 6, 7, 8, 4, 7, 2, 1, 8, 4, 2, 3, 9]
Constructed Result List : [6, 7, 7, 2, 2, 3]

Time complexity: O(n), where n is the length of the input list.

Auxiliary space: O(n), where n is the length of the input list.