Python|将一维列表转换为可变长度的二维列表

给定一维列表“lst”和可变长度列表“var_lst”，编写一个Python程序将给定的一维列表转换为给定可变长度的二维列表。

例子：

Input : lst = [1, 2, 3, 4, 5, 6]
        var_lst = [1, 2, 3]
Output : [[1], [2, 3], [4, 5, 6]]

Input : lst = ['a', 'b', 'c', 'd', 'e']
        var_lst = [3, 2]
Output : [['a', 'b', 'c'], ['d', 'e']]

方法#1：列表切片

# Python3 program to Convert 1D 
# list to 2D list
from itertools import islice
  
def convert(lst, var_lst):
    idx = 0
    for var_len in var_lst:
        yield lst[idx : idx + var_len]
        idx += var_len
      
# Driver code
lst = [1, 2, 3, 4, 5, 6]
var_lst = [1, 2, 3]
print(list(convert(lst, var_lst)))

输出：

[[1], [2, 3], [4, 5, 6]]

方法 #2：使用itertools.islice()

# Python3 program to Convert 1D 
# list to 2D list
from itertools import islice
  
def convert(lst, var_lst):
    it = iter(lst)
    return [list(islice(it, i)) for i in var_lst]
      
# Driver code
lst = [1, 2, 3, 4, 5, 6]
var_lst = [1, 2, 3]
print(convert(lst, var_lst))

输出：

[[1], [2, 3], [4, 5, 6]]