Как ускорить этот код на Python?

Question

Задача: В конце занятия n студентов попросили у Кирилла задачи, которых у него всего m. Каждый ученик хотел бы получить a задач, и если он получит меньше, а именно b, то будет зол на (a −b)^2 злости. Каждую задачу можно дать не более чем одному ученику. Кирилл очень не хотел бы злить своих любимых студентов, поэтому просит вас помочь минимизировать суммарную злость в классе.

def angry(m, n, List):
    numb = 0
    for i in range(m):
        List.sort(reverse=True)
        List[0] -= 1
    for i in range(n):
        numb += List[i] ** 2 
    return numb

m, n = map(int, input().split())
a = list(map(int, input().split()))

# Вывод минимальной злости
print(angry(m, n, a) % (2 ** 64))

также пробовал

def angry(m, n, List):
    for i in range(m):
        List.sort(reverse=True)
        List[0] -= 1
    return sum([x ** 2 for x in List])

m, n = map(int, input().split())
a = list(map(int, input().split()))

# Вывод минимальной злости
print(angry(m, n, a) % (2**64))

Ошибка "Превышено максимальное время работы".

Как можно ускорить этот код, чтобы конечный результат был таким же?

Answer 1

Для решения задачи можно использовать кучу (heap). Куча позволяет в эффективное время извлекать максимальный элемент и затем возвращать его обратно в кучу с уменьшенным значением. Это уменьшит время работы до O(m log n) вместо O(m * n log n).

P.S. Если оставить ваш код как есть, его трудно будет оптимизировать, поэтому нужно изменить алгоритм.

Пример:

import heapq

def angry(m, n, tasks):
    # Превратим список задач в max-кучу через использование отрицательных значений
    heap = [-x for x in tasks]
    heapq.heapify(heap)
    
    # Раздаем задачи
    for _ in range(m):
        # Извлекаем наибольшее количество задач
        max_tasks = heapq.heappop(heap)
        # Уменьшаем его на 1 и кладем обратно в кучу
        heapq.heappush(heap, max_tasks + 1)  # т.к. max_tasks отрицательное, прибавляем 1

    # Считаем итоговую злость
    total_anger = sum(x ** 2 for x in heap)
    
    # Возвращаем результат с нужным модулем
    return -total_anger % (2 ** 64)

# Чтение входных данных
m, n = map(int, input().split())
tasks = list(map(int, input().split()))

# Вывод минимальной злости
print(angry(m, n, tasks))

Answer 2

Не уверен, но на всякий случай:

import os

print("-" * 50 + "\nСуммарная злость в классе:\n" + "-" * 50)

def minimize_anger(n, m, a):
 # Если задач меньше, чем студентов, то невозможно удовлетворить всех:
 if m < n: return "Невозможно удовлетворить всех студентов"

 # Изначально все студенты получают по одной задаче:
 tasks_per_student = [1] * n
 remaining_tasks = m - n

 # Распределяем оставшиеся задачи:
 for i in range(n):
  if remaining_tasks == 0: break
  tasks_per_student[i] += 1
  remaining_tasks -= 1

 # Вычисляем суммарную злость:
 total_anger = sum((a - tasks_per_student[i]) ** 2 for i in range(n))
 return total_anger

# Пример использования:
n = 5  # количество студентов
m = 10  # количество задач
a = 3  # желаемое количество задач для каждого студента

print("Количество студентов:", n)
print("Количество задач:", m)
print("Желаемое количество задач для каждого студента:", a)
print("Минимизированная злость в классе:", minimize_anger(n, m, a))

print("\nНажмите любую клавишу для продолжения...")
os.system("pause > nul" if os.name == "nt" else "read > /dev/null")

Answer 3

В ограничениях число задач может быть до двух миллиардов. Раздавать их по одной никакого времени не хватит.

Если два ученика получили задачи и их остаточная злобность (a_i - b_i) различается больше чем единицу, то можно одну задачу передать от менее злобного более злобному и общая злобность уменьшится. Откуда решение должно иметь такой вид: кто-то не получает задач вовсе, а те кто их получает, получает так чтобы злобность была как можно равномернее распределена. С этой точки зрения решение в вопросе верное: выбрать самого злобного, дать ему задачу, обновить его злобность и повторять пока задачи не закончатся или пока злобных не останется. Но как уже сказано, раздача по одной задаче за раз - тупик.

Пусть три ученика попросили по шесть задач, два ученика по четыре и два по три. Число задач m = 21. Отсортируем запросы:

        * * *
        * * *
    * * * * *     m = 21
* * * * * * *
* * * * * * *
* * * * * * *

Разобъём их на группы по высоте:

        3 3 3
        3 3 3
    2 2 3 3 3     m = 21
1 1 2 2 3 3 3
1 1 2 2 3 3 3
1 1 2 2 3 3 3

Нужно срезать самую высокую группу до высоты второй самой высокой. Для этого нужно шесть задач. Они у нас есть:

    2 2 3 3 3
1 1 2 2 3 3 3     m = 15
1 1 2 2 3 3 3
1 1 2 2 3 3 3

Группы одинаковой высоты объединим:

    2 2 2 2 2
1 1 2 2 2 2 2     m = 15
1 1 2 2 2 2 2
1 1 2 2 2 2 2

Снова хотим срезать самую высокую группу. Требуется пять задач. Они есть. Срезаем и объединяем группы:

1 1 1 1 1 1 1
1 1 1 1 1 1 1
1 1 1 1 1 1 1     m = 10

На то чтобы срезать эту группу не хватает задач. Уберём столько рядов сколько сможем (один) и из одного ряда удалим оставшиеся задачи (три):

      1 1 1 1
1 1 1 1 1 1 1     m = 0

Таких картинок как выше не бывает, в одной группе все должны иметь одну высоту. Надо сделать две группы:

      2 2 2 2
1 1 1 2 2 2 2     m = 0

Так как задач не осталось, останавливаем цикл и считаем оставшуюся злость.

Сложность программы O(n·log n).

В список исходных злобностей добавлен один ученик с нулевыми запросами. Он нужен, чтобы обработка последней оставшейся группы не отличалась от остальных.

import itertools

m, n = map(int, input().split())
a = list(map(int, input().split()))
a.append(0)

stack = list((a, sum(1 for _ in g)) for a, g in itertools.groupby(sorted(a)))

while m > 0 and len(stack) >= 2:
    a1, w1 = stack[-1]
    a2, w2 = stack[-2]
    h = a1 - a2
    if m >= h * w1:
        stack.pop()
        stack.pop()
        stack.append((a2, w1 + w2))
        m -= h * w1
    else:
        q, r = divmod(m, w1)
        stack.pop()
        stack.append((a1 - q - 1, r     ))
        stack.append((a1 - q    , w1 - r))
        m = 0

print(sum(a * a * w for a, w in stack) % 2 ** 64)

БЛОГ НА HUSL

Как ускорить этот код на Python?

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