Оптимизация сортировки - Python

У меня есть алгоритм по сортировке большого массива данных. Он сильно нагружает систему, и во время запуска сортировщика, программа зависает (Но позже возвращается в нормальную работу). Сам процесс работает нормально и выдаёт нужные мне данные.

Есть массив словарей типа {"id": 1, "group": 2, "sub_group": 3} Есть 2 вида сортировки: default и by_id, а так же список с двумя вложенными списками [[], []], назовём его array

--- default ---

Если у словаря нет группы, значит он отправляется в array[0], где должен встать в массив со своей подгруппой. В итоге должна получится такая иерархия: array[0] -> sub_group. Если же у словаря есть группа, то должен добавится массив с группой в array[1] в свою подгруппу. Тогда, чтобы достать элемент, нужно проделать такой путь: array[1][group][sub_group]

--- by_id ---

Здесь принцип абсолютно такой-же, но все эти элементы должны сортироваться по id Тогда, чтобы добраться до элемента, нужно проделать следующий путь: array[id][1][group][sub_group] или же array[id][0][sub_group]

На данный момент у меня такой вот бардак и полная неразбериха в коде

import random


def data_generator():
    data = []
    for i in range(random.randint(1200, 1800)):
        site = {
            "user_id": random.randint(1, 150),
            "group": random.randint(1, 5),
            "sub_group": random.randint(1, 20)
        }
        data.append(site)
    return data


def sorter(web_sites, sort_type):
    if sort_type == 0 or sort_type == 2:  # default
        sites = [[], []]
        groups = []
        sub_groups = [[]]
        unsorted = []

        for site in web_sites:
            group = site['group']

            if group:
                sub_group = site['sub_group']

                if group not in groups:
                    groups.append(group)
                    sites[1].append([])
                    sub_groups.append([])

                g_index = groups.index(group)

                if sub_group not in sub_groups[g_index]:
                    sub_groups[g_index].append(sub_group)
                    sites[1][g_index].append([])
                sites[1][g_index][sub_groups[g_index].index(sub_group)].append(site)

            else:
                sub_group = site['sub_group']

                if sub_group not in unsorted:
                    unsorted.append(sub_group)
                    sites[0].append([])

                sites[0][unsorted.index(sub_group)].append(site)
        return sites

    if sort_type == 1:  # by user
        sites = []
        ids = []
        groups = [[]]
        sub_groups = [[[]]]
        unsorted = [[]]

        for site in web_sites:
            group = site['group']  # группа
            user_id = site['user_id']  # подгруппа

            if user_id not in ids:  # проверяем, есть ли в списках id текущий
                ids.append(user_id)
                sites.append([[], []])
                groups.append([])
                sub_groups.append([[]])
                unsorted.append([])
            ui_index = ids.index(user_id)

            if group:  # проверяем, есть ли группа у элемента
                sub_group = site['sub_group']

                if group not in groups[ui_index]:  # проверяем, есть ли в группах текущая
                    groups[ui_index].append(group)
                    sites[ui_index][1].append([])
                    sub_groups[ui_index].append([])
                g_index = groups[ui_index].index(group)

                if sub_group not in sub_groups[ui_index][g_index]:
                    sub_groups[ui_index][g_index].append(sub_group)
                    sites[ui_index][1][g_index].append([])
                sites[ui_index][1][g_index][sub_groups[ui_index][g_index].index(sub_group)].append(site)

            else:
                sub_group = site['sub_group']  # подгруппа - это домен сайта

                if sub_group not in unsorted[ui_index]:
                    unsorted[ui_index].append(sub_group)
                    sites[ui_index][0].append([])

                sg_index = unsorted[ui_index].index(sub_group)
                sites[ui_index][0][sg_index].append(site)
        return sites

Ответы (1 шт):

Автор решения: Stanislav Volodarskiy

Работу замедляют проверки вида not in для списков и вызовы .index() для них же. Если рядом со списком завести словарь, то оба вызова из линейных (от длины списка) станут константными (обращение к словарю). Кроме того, следует избавиться от дублирования кода.

Функция locate строит таблицу из вложенных списков и возвращает самый вложенный список. Аргумент tree - пара из словаря и списка верхнего уровня.

Функция протестирована на данных из вопроса. Ускорение в 1-2 раза.

def locate(tree, *route):
    index, table = tree
    for i in route: 
        if i not in index:
            lst = []
            index[i] = {}, lst
            table.append(lst)
        index, table = index[i]
    return table


def sorter2(web_sites, sort_type):
    if sort_type == 0 or sort_type == 2:  # default
        gtree = {}, []
        utree = {}, []

        for site in web_sites:
            group = site['group']
            if group:
                table = locate(gtree, group, site['sub_group'])
            else:
                assert False, 'not tested'
                table = locate(utree, site['sub_group'])
            table.append(site)

        return [utree[1], gtree[1]]

    if sort_type == 1:  # by user
        tree = {}, []

        for site in web_sites:
            pair = locate(tree, site['user_id'])
            if not pair:
                pair.extend((({}, []), ({}, [])))

            group = site['group']  # группа
            if group:  # проверяем, есть ли группа у элемента
                table = locate(pair[1], group, site['sub_group'])
            else:
                assert False, 'not tested'
                table = locate(pair[0], site['sub_group'])
            table.append(site)

        return [[t1, t2] for (_, t1), (_, t2) in tree[1]]
→ Ссылка