How to write a function that does the following -

After each occurrence of num1 in a list, insert num2

def insert_after(lst, num1, num2):

CodePudding user response：

In general, people on stack overflow would appreciate you to at least attempt a solution before asking for help. That being said, here is an O(1) solution that is the first to come to mind:

PREVIOUS ANSWER WAS INCORRECT

The best way to do this is as follow: Iterating through the length of the list as I did before does not work, as the list is growing when you insert but you are iterating through a fixed number.

def insert_after(lst,num1,num2):
    output = []
    for num in lst:
        output.append(num)
        if num == num1:
            output.append(num2)
    return output

CodePudding user response：

Here are a couple fun ways using list comprehensions:

Solution 1:

lst = [1,5,4,3,5,5,7,8,9,6,5]

def insert_after_more_readable(lst,num1,num2):
    temp_lst = [[num] if num != num1 else [num, num2] for num in lst]
    final_lst = [elem for sublist in temp_lst for elem in sublist]
    return final_lst
print(lst, 5,8)

Explanation:

1. We create a list of lists called temp_lst. Every time we encounter num1, we insert a sublist [num1,num2], and otherwise we insert sublist [num].
2. We iterate over all the sublists, and take out all their elements and put them into the outer list.

Solution 2:

This is a less readable version, but the exact same thing using nested list comprehensions, all in one line.

lst = [1,5,4,3,5,5,7,8,9,6,5]
def insert_after(lst, num1, num2):
    return [elem for sublist in [[num] if num != num1 else [num,num2] 
for num in lst] for elem in sublist]
print(lst, 5,8)

In this version, the inner list that is formed is equivalent to temp_lst above.