【python】将数组list等分成几组

写在前面

问题描述：
给定一个数组，将其等分成N组，最后一组元素不足也作为一组。
例如：将数组test_list = [3, 4, 5, 6, 7, 22, 35, 46, 78, 100]，等分成数组：alist = [[3, 4, 5], [6, 7, 22], [35, 46, 78], [100]]

对于上述问题的反向操作，也就是，给定数组alist，将其合并合并成一个数组test_list，python的内置库itertools的chain函数可实现改操作。

import itertools
alist = [[3, 4, 5], [6, 7, 22], [35, 46, 78], [100]]
print(list(itertools.chain(alist[0], alist[1])))
print(list(itertools.chain(*alist)))  # 参数化

len_a = len(alist)
new_list = []
for i in range(len_a):
    new_list = list(itertools.chain(new_list, alist[i]))

print(new_list)

[3, 4, 5, 6, 7, 22]
[3, 4, 5, 6, 7, 22, 35, 46, 78, 100]
[3, 4, 5, 6, 7, 22, 35, 46, 78, 100]

实现代码

下面是实现上述问题的代码。若将数组等分成n组，每组元素至少m个，其中 m = len_list / n的上取整(用到math模块中的ceil)。因此，该数组每当下标整除m时，在新数组中增加一个数组，不能整除时，就在增加的数组后添加元素。

def lst_trans0_new(test_list, n):
    """n: split list into `n` groups
    """
    len_list = len(test_list)
    m_ceil = math.ceil(len_list / float(n))
    m_floor = math.floor(len_list / float(n))
    if len_list / m_ceil < n:
        m = m_floor
    else:
        m = m_ceil
    alist = []
    group_m = -1
    for i in range(len_list):
        if i % m == 0:
            if group_m == n - 2:
                alist.append(test_list[i:])
                break
            else:
                group_m += 1
                alist.append([test_list[i]])
        else:
            alist[group_m].append(test_list[i])
    return alist

# 使用：
test_list = [3, 4, 5, 6, 7, 22, 35, 46, 78, 100, 2, 3]  # , 4]
print(len(test_list))  ＃ 12
print(lst_trans0_new(test_list, 5))
# [[3, 4], [5, 6], [7, 22], [35, 46], [78, 100, 2, 3]]

以上为修正后的代码。

---

原来错误代码：【忽略】 统一按 ceil(len / n)上整：得到每个子元素列表中的个数。返回的划分列表数 确实可能和n不等。

In [2]: import math
   ...: test_list = [3, 4, 5, 6, 7, 22, 35, 46, 78, 100]
   ...: len_list = len(test_list)
   ...: n = 3
   ...: m = math.ceil(len_list / float(n))
   ...: alist = []
   ...: group_m = -1
   ...: for i in range(len_list):
   ...:     if i % m == 0:
   ...:         group_m += 1
   ...:         alist.append([test_list[i]])
   ...:         print "%s: %s" % (i, alist)
   ...:     else:
   ...:         alist[group_m].append(test_list[i])
   ...:         print "%s: %s" % (i, alist)
   ...: print "split list:", alist
   ...: 
0: [[3]]
1: [[3, 4]]
2: [[3, 4, 5]]
3: [[3, 4, 5, 6]]
4: [[3, 4, 5, 6], [7]]
5: [[3, 4, 5, 6], [7, 22]]
6: [[3, 4, 5, 6], [7, 22, 35]]
7: [[3, 4, 5, 6], [7, 22, 35, 46]]
8: [[3, 4, 5, 6], [7, 22, 35, 46], [78]]
9: [[3, 4, 5, 6], [7, 22, 35, 46], [78, 100]]
split list: [[3, 4, 5, 6], [7, 22, 35, 46], [78, 100]]

# test_split.py
#!/usr/bin/env python
# -*- coding: utf-8 -*-
import math

def lst_trans0(test_list, n):
    """n: split list into `n` groups
    """
    len_list = len(test_list)
    m = math.ceil(len_list / float(n))
    alist = []
    group_m = -1
    for i in range(len_list):
        if i % m == 0:
            group_m += 1
            alist.append([test_list[i]])
        else:
            alist[group_m].append(test_list[i])
    return alist

def lst_trans1(lst, n):
    m = int(math.ceil(len(lst) / float(n)))
    sp_lst = []
    for i in range(n):
        sp_lst.append(lst[i*m:(i+1)*m])
    return sp_lst

# 可以不使用math模块
def lst_trans2(lst, n):
    if len(lst) % n != 0:
        m = (len(lst) // n) + 1
    else:
        m = len(lst) // n
    sp_lst = []
    for i in range(n):
        sp_lst.append(lst[i*m:(i+1)*m])
    return sp_lst

if __name__ == '__main__':
    test_list = [3, 4, 5, 6, 7, 22, 35, 46, 78, 100]
    print(lst_trans0(test_list, 3))
    print(lst_trans1(test_list, 3))
    print(lst_trans2(test_list, 3))

执行结果：

$ python test_split.py
[[3, 4, 5, 6], [7, 22, 35, 46], [78, 100]]
[[3, 4, 5, 6], [7, 22, 35, 46], [78, 100]]
[[3, 4, 5, 6], [7, 22, 35, 46], [78, 100]]

---

(完)