Search code examples
pythonlist-comprehension

Multiple loop in list comprehension


I have a problem set in hand, the description is something like:

We have 2 lists ie "nums1" and "nums2" and 2 integers m and n

nums1 = [1, 2, 3, 7, 8, 9]
m = 3
nums2 = [2, 5, 6]
n = 2

Here "m" denotes number of elements to be picked from "nums1" and "n" denotes the number of elements to be picked from "nums2". There are other ways ofcourse, but I want to create a list with only "list comprehension" on the above conditions.

With a simple for loop I can achieve that with:

li = []
for i in range(m):
  li.append(nums1[i])
for j in range(n):
  li.append(nums2[j])
print(li)

I tried the following but the outcome result is not as expected:

[a for a in nums1 for b in nums2]

Expected output: [1, 2, 2, 3, 5]

Also, does list comprehension supports multiple loop?


Solution

  • Given the following:

    nums1 = [1, 2, 3, 7, 8, 9]
    m = 3
    nums2 = [2, 5, 6]
    n = 2
    

    Using strictly list comprehension, the following does what you asked for:

    sorted([nums1[a] for a in range(m)] + [nums2[b] for b in range(n)])
    

    Returns:

    [1, 2, 2, 3, 5]
    

    as per the conditions with "List comprehension" to be specific.

    Although as others have said in the comment, a slice would have been more pythonic and easier to reason about. But since you asked for a list comprehension implementation, then this is as close to an answer as you'll get.