Как можно оптимизировать по времени мой код?

Question

Суть задачи: вводятся числа n,m (0<=n,m<=10^5) - размеры матрицы, затем сама матрица в n строках. Также вводятся число q (1<=q<=10^5) количество запросов. В последующих q строках вводятся по три числа i,j,k. i,j (0<=i<n,0<=j<m) - номера столбца и строки, которые нужно выцепить из матрицы, их данные отсортировать (значение на пересечении берется 1 раз), и вывести k-й(1<=k<=m+n-1) элемент среди них.

Примечание: Для первого запроса: 3 2 1 построим и отсортируем объединенный массив: [3 22 55 64 84 87]. Элемент на 1 позиции - 3. Для третьего запроса: 3 2 6 аналогичный массив, как для первого. Элемент на 6 позиции - 87. Для четвертого запроса: 0 1 4 построим и отсортируем объединенный массив: [8 22 30 55 76 84]. Элемент на 4 позиции - 55.

Вот мой код:

import random

def qsort(a):
    if len(a)<=1: return a
    x=random.choice(a)
    bl=[b for b in a if b<x]
    bx=[b for b in a if b==x]
    br=[b for b in a if b>x]
    return qsort(bl)+bx+qsort(br)
n,m=list(map(int,input().split()))
mat=[list(map(int,input().split())) for i in range(n)]
mat1=list(zip(*mat))
q=int(input())
for i in range(q):
     k,l,m=list(map(int,input().split()))
     s1=list(mat1[l])
     s1=s1[:k]+s1[k+1:]
     print(qsort(s1+mat[k])[m-1])

Я всё равно спотыкаюсь об ограничения в 1.5 секунды. Мой код выполняется за 1.586 из-за чего и получает TL. Переводил в С++ всё равно не проходит барьер по времени. Как можно оптимизировать по времени данный код, или есть ещё какие-нибудь идеи?

Answer 1

Самый быстрый вариант, который пока смог получить. Есть версия на алгоритме Quickselect, но она оказалась в два раза медленнее стандартной Питоновской сортировки. Пока не понял почему. Подозреваю, что Питоновская сортировка реализована на C, потому тупо быстрее алгоритма Quickselect, реализованного на Питоне. В итоге оставил версию с сортировкой.

P.S. Важно входные данные matrix и row_col_key оформить в виде кортежа. При передаче в функцию get_matrix_element, кортежи передаются как копии, а списки - как ссылки. При работе с ссылками возникают затраты времени на поиск данных в памяти по ссылке. С кортежами, как с локальными данными, функция работает чуть быстрее.

P.S.S. Пробовал реализовать кеш на словаре, но на таких маленьких данных и при незначительной повторяемости, затраты на содержание словаря оказались выше.

import itertools

matrix = (
    (13, 84, 22),
    (78, 76, 64),
    (80, 8, 87),
    (3, 55, 84),
)
row_col_key = (
    (3, 2, 1),
    (3, 0, 5),
    (3, 2, 6),
    (0, 1, 4),
    (1, 2, 3),
    (2, 2, 4),
    (1, 2, 6),
    (1, 1, 6),
)

def get_matrix_element(matrix, row_col_key):
    matrix_transp = list(zip(*matrix))
    for r, c, k in row_col_key:
        yield sorted(
            itertools.chain(matrix[r][:c], matrix[r][c + 1 :], matrix_transp[c])
        )[k - 1]

print(list(get_matrix_element(matrix, row_col_key)))

Ниже версия на Quickselect:

def _get_index(data, first, last): 
    base_element = data[last]
    index = first
    for i in range(first, last):
        if data[i] <= base_element:
            if index != i:
                data[index], data[i] = data[i], data[index]
            index += 1
    if index != last:
        data[index], data[last] = data[last], data[index]
    return index


def smallest_in_position(data, position):
    first = 0
    last = len(data) - 1
    iposition = position - 1
    if (position > first and position <= last + 1):
        if first == iposition:
            return min(data)
        if last == iposition:
            return max(data)
        while (index := _get_index(data, first, last)) != iposition:
            if (index > iposition):
                last = index - 1
            else:
                first = index + 1
        return data[index]
    else:
        raise IndexError("Position out of range")

def get_matrix_element2(matrix, row_col_key):
    matrix_transp = list(zip(*matrix))
    for r, c, k in row_col_key:
        _row_plus_col = list(itertools.chain(matrix[r][:c], matrix[r][c + 1:], matrix_transp[c]))
        yield smallest_in_position(_row_plus_col, k)

Версия без сортировки на каждой итерации цикла. Матрица сортируется единожды до входа в цикл. По скорости практически равна обычной Питоновской сортировке (см. самую первую версию). Не уверен, что эта версия пройдет тест (знать бы еще как тест работает).

def get_matrix_element_3(matrix, row_col_key):
    # Предварительная единоразовая подготовка данных. Можно вынести за пределы функции
    # Формируем словарь со всеми комбинациями отсортированных строк без общего элемента со столбцами
    rows_sorted = dict(
        ((irow, icol), sorted(val for i, val in enumerate(matrix[irow]) if i != icol))
        for icol in range(len(matrix[0]))
        for irow in range(len(matrix))
    )
    # Транспонируем матрицу и сортируем. Получаем отсортированные столбцы
    cols_sorted = tuple(map(sorted, zip(*matrix)))
    row_size = len(matrix[0]) - 1
    col_size = len(matrix)
    # Запускаем цикл поиска k-го элемента
    for r, c, k in row_col_key:
        row = rows_sorted[(r,c)]
        col = cols_sorted[c]
        # Сначала отрабатываем крайние случаи
        if k == 1:
            yield row[0] if row[0] < col[0] else col[0]
        elif k == row_size + col_size:
            yield row[-1] if row[-1] > col[-1] else col[-1]
        else:
            # Шагаем по двум спискам и отбираем наименьшие значения
            irow = icol = 0
            while k > 0 and irow < row_size and icol < col_size:
                if col[icol] < row[irow]:
                    element = col[icol]
                    icol += 1
                else:
                    element = row[irow]
                    irow += 1
                k -= 1
            
            if k > 0:
                # Если один из списков закончился раньше, чем достигнуто k
                yield row[irow + k - 1] if irow < row_size < 0 else col[icol + k - 1]
            else:
                yield element

БЛОГ НА HUSL

Как можно оптимизировать по времени мой код?

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