1d преобразование массива: распределение групп разных размеров в уникальные партии с определенными условиями

Question

Проблема: Создайте наиболее эффективную функцию, чтобы превратить 1d массив (столбец group_id) в другой 1d массив (выходной столбец).

Условия:

1. Максимум n групп может быть в любой партии, в этом примере n=2.

2. Каждая партия должна содержать группы одинакового размера.

3. Тривиальное условие: свести к минимуму количество партий.

Функция будет распределять эти группы разного размера по партиям с уникальными идентификаторами при условии, что каждая партияимеет фиксированный размер И каждая партия содержит только группы с одинаковым размером.

data = {'group_size': [1,2,3,1,2,3,4,5,1,2,1,1,1],
        'batch_id':   [1,4,6,1,4,6,7,8,2,5,2,3,3]}
df = pd.DataFrame(data=data)
print(df)

    group_size  batch_id
0          1       1
1          2       4
2          3       6
3          1       1
4          2       4
5          3       6
6          4       7
7          5       8
8          1       2
9          2       5
10         1       2
11         1       3
12         1       3

Что мне нужно:

some_function( data['group_size'] ), чтобы дать мне data['batch_id']

Редактировать:

Моя неуклюжая функция

def generate_array():

    out = 1
    batch_size = 2
    dictionary = {}

    for i in range(df['group_size'].max()):
        # get the mini df corresponding to the group size
        sub_df = df[df['group_size'] == i+1 ]
        # how many batches will we create?
        no_of_new_batches = np.ceil ( sub_df.shape[0] / batch_size )
        # create new array
        a = np.repeat(np.arange(out, out+no_of_new_batches ), batch_size)
        shift = len(a) - sub_df.shape[0]

        # remove last elements from array to match the size
        if len(a) != sub_df.shape[0]:
            a = a[0:-shift]

        # update batch id
        out = out + no_of_new_batches

        # create dictionary to store idx
        indexes = sub_df.index.values

        d = dict(zip(indexes, a))

        dictionary.update(d)

    array = [dictionary[i] for i in range(len(dictionary))]

    return array

generate_array()
Out[78]:
[1.0, 4.0, 6.0, 1.0, 4.0, 6.0, 7.0, 8.0, 2.0, 5.0, 2.0, 3.0, 3.0]

Answer 1

Вот мое решение.Я не думаю, что он дает точно такой же результат, как ваша функция, но он удовлетворяет вашим трем правилам:

import numpy as np

def package(data, mxsz):
    idx = data.argsort()
    ds = data[idx]
    chng = np.empty((ds.size + 1,), bool)
    chng[0] = True
    chng[-1] = True
    chng[1:-1] = ds[1:] != ds[:-1]
    szs = np.diff(*np.where(chng))
    corr = (-szs) % mxsz
    result = np.empty_like(idx)
    result[idx] = (np.arange(idx.size) + corr.cumsum().repeat(szs)) // mxsz
    return result

data = np.random.randint(0, 4, (20,))
result = package(data, 3)
print(f'group_size {data}')
print(f'batch_id   {result}')
check = np.lexsort((data, result))
print('sorted:')
print(f'group_size {data[check]}')
print(f'batch_id   {result[check]}')

Пример запуска с n = 3, последние две строки вывода совпадают спервые два, отсортированы только для упрощения проверки:

group_size [1 1 0 1 2 0 2 2 2 3 1 2 3 2 1 0 1 0 2 0]
batch_id   [3 3 1 3 6 1 6 5 6 7 2 5 7 5 2 1 2 0 4 0]
sorted:
group_size [0 0 0 0 0 1 1 1 1 1 1 2 2 2 2 2 2 2 3 3]
batch_id   [0 0 1 1 1 2 2 2 3 3 3 4 5 5 5 6 6 6 7 7]

Как это работает:

1) сортировка данных

2) обнаружение, где изменяются отсортированные данные, чтобы идентифицировать группы равныхзначения («группы размеров групп»)

3) определяют размеры групп размеров групп и для каждого вычисления, что пропускает чистое кратное n

4), перечисляют отсортированные данные, покапри каждом переключении на новую группу размеров групп, переходящих к следующему чистому кратному n;мы используем (3), чтобы сделать это в векторизованном виде

5) деление по полу на n, чтобы получить идентификаторы партии

6) перемешать обратно в исходный порядок