python - 分区组作为 Python 和 NumPy 中的矩阵列表

Question

我有一个以下形式的 NumPy 矩阵:

dummy = np.array([['A', 1, 1],
                  ['A', 1, 1],
                  ['B', 1, 1],
                  ['C', 1, 1],
                  ['F', 1, 1],
                  ['I', 1, 1],
                  ['I', 1, 1],
                  ['I', 1, 1],
                  ['J', 1, 1],
                  ['K', 1, 1],
                  ['L', 1, 1],
                  ['M', 1, 1]])

我试图生成的是一个 NumPy 矩阵列表，其中每个矩阵只能有 n 个不同的字符串值变体，例如:

n = 4:

[array([['A', '1', '1'],
        ['A', '1', '1'],
        ['B', '1', '1'],
        ['C', '1', '1'],
        ['F', '1', '1']], <= 4 different groups (A,B,C,F)
       dtype='|S1'),
 array([['I', '1', '1'],
        ['I', '1', '1'],
        ['I', '1', '1'],
        ['J', '1', '1'],
        ['K', '1', '1'],
        ['L', '1', '1']], <= 4 different groups (I,J,K,L)
       dtype='|S1'),
 array([['M', '1', '1']], <= only one but kept
       dtype='|S1')]

我有以下几乎可以工作的功能..

def partition_by(x, groups):
    uniques = set([])
    p = []
    q = []
    for i in x:
        if len(uniques) < groups or i[0] in uniques:
           uniques.add(i[0])
           p.append(i.tolist())
    q.append(np.array(p))
    return q

partition_by(dummy, 4)

产量:

[array([['A', '1', '1'],
       ['A', '1', '1'],
       ['B', '1', '1'],
       ['C', '1', '1'],
       ['F', '1', '1']], 
      dtype='|S1')]

Answer 1

这是一种方法-

def split_col_based(dummy, colID = 0, n=4):
    mask = dummy[1:,colID] != dummy[:-1,colID]
    interval_idx = n*(np.arange((mask.sum()+1)//n)+1)
    idx = np.searchsorted(mask.cumsum(), interval_idx)
    return np.split(dummy, idx+1, axis=0)

示例输入、输出

1) 输入数组:

In [79]: dummy
Out[79]: 
array([['A', '1', '1'],
       ['A', '1', '1'],
       ['B', '1', '1'],
       ['C', '1', '1'],
       ['F', '1', '1'],
       ['I', '1', '1'],
       ['I', '1', '1'],
       ['I', '1', '1'],
       ['J', '1', '1'],
       ['K', '1', '1'],
       ['L', '1', '1'],
       ['M', '1', '1']], 
      dtype='|S1')

2) n=4 输出:

In [80]: split_col_based(dummy, n=4)
Out[80]: 
[array([['A', '1', '1'],
        ['A', '1', '1'],
        ['B', '1', '1'],
        ['C', '1', '1'],
        ['F', '1', '1']], 
       dtype='|S1'), array([['I', '1', '1'],
        ['I', '1', '1'],
        ['I', '1', '1'],
        ['J', '1', '1'],
        ['K', '1', '1'],
        ['L', '1', '1']], 
       dtype='|S1'), array([['M', '1', '1']], 
       dtype='|S1')]

3) n=5 输出:

In [81]: split_col_based(dummy, n=5)
Out[81]: 
[array([['A', '1', '1'],
        ['A', '1', '1'],
        ['B', '1', '1'],
        ['C', '1', '1'],
        ['F', '1', '1'],
        ['I', '1', '1'],
        ['I', '1', '1'],
        ['I', '1', '1']], 
       dtype='|S1'), array([['J', '1', '1'],
        ['K', '1', '1'],
        ['L', '1', '1'],
        ['M', '1', '1']], 
       dtype='|S1')]

4) n=2 输出:

In [84]: split_col_based(dummy, n=2)
Out[84]: 
[array([['A', '1', '1'],
        ['A', '1', '1'],
        ['B', '1', '1']], 
       dtype='|S1'), array([['C', '1', '1'],
        ['F', '1', '1']], 
       dtype='|S1'), array([['I', '1', '1'],
        ['I', '1', '1'],
        ['I', '1', '1'],
        ['J', '1', '1']], 
       dtype='|S1'), array([['K', '1', '1'],
        ['L', '1', '1']], 
       dtype='|S1'), array([['M', '1', '1']], 
       dtype='|S1')]