Оптимизация кода с бинарным поиском

По этой задаче был уже задан вопрос по оптимизации, но можно ли ещё оптимизировать код, не подключая никаких библиотек, поработав с этой частью кода?

while j + c - 1 < len(mass):
    if mass[j + c - 1] - mass[j] <= m:
        k += 1
        j = j + c
    else:
        j = j + 1

Задача №1620. Субботник

В классе учатся N человек. Классный руководитель получил указание направить на субботник R бригад по С человек в каждой. Все бригады на субботнике будут заниматься переноской бревен. Каждое бревно одновременно несут все члены одной бригады. При этом бревно нести тем удобнее, чем менее различается рост членов этой бригады.

Числом неудобства бригады будем называть разность между ростом самого высокого и ростом самого низкого членов этой бригады (если в бригаде только один человек, то эта разница равна 0). Классный руководитель решил сформировать бригады так, чтобы максимальное из чисел неудобства сформированных бригад было минимально. Помогите ему в этом!

Рассмотрим следующий пример:

Пусть в классе 8 человек, рост которых в сантиметрах равен 170, 205, 225, 190, 260, 130, 225, 160, и необходимо сформировать две бригады по три человека в каждой.

Тогда одним из вариантов является такой:

1 бригада: люди с ростом 225, 205, 225

2 бригада: люди с ростом 160, 190, 170

При этом число неудобства первой бригады будет равно 20, а число неудобства второй — 30. Максимальное из чисел неудобств будет 30, и это будет наилучший возможный результат.

Формат входных данных

Сначала вводятся натуральные числа N, R и C — количество человек в классе, количество бригад и количество человек в каждой бригаде (1 ≤ R∙C ≤ N ≤ 100 000). Далее вводятся N целых чисел — рост каждого из N учеников. Рост ученика — натуральное число, не превышающее 1 000 000 000.

Формат выходных данных

Выведите одно число — наименьше возможное значение максимального числа неудобства сформированных бригад.

n, r, c = map(int, str(input()).split())
mass = []
for i in range(n):
    mass.append(int(input()))
mass = sorted(mass)
right = mass[len(mass) - 1] - mass[0]
left = 0
while left < right:
    m = (left + right) // 2
    j = 0
    k = 0
    while j + c - 1 < len(mass):
        if mass[j + c - 1] - mass[j] <= m:
            k += 1
            j = j + c
        else:
            j = j + 1
    if k >= r:
        right = m

    else:
        left = m + 1

print(right)

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

Автор решения: Vladimir Bogdanov

Создаем словарь неудобств с суммой первого и последнего индекса группы в качестве ключей. Находим ключ с минимальным неудобством. Следующие минимальные неудобства ищем в непересекающихся группах.

r, c = 2, 3
mass = [170, 205, 225, 190, 260, 130, 225, 160]
mass.sort()
inconv = {a + b: mass[b] - mass[a] for a, b in enumerate(range(c - 1, len(mass)))}
offset: int = 2 * c * (r - 1)
#Список ключей, отсортированных по значениям словаря.
sorted_keys = sorted(inconv.keys(), key=lambda _key: inconv[_key])
result = None
for start_key in sorted_keys:
    try:
        #Одношаговый генератор сразу находит нужный ответ
        result_key=next(
            key for key in sorted_keys
            if abs(key - start_key) >= offset
        )
    except StopIteration:
        #Если поиск потерпел неудачу, переходим к следующему ключу из sorted_keys
        continue
    else:
        #Успешный поиск изымает значение из словаря и досрочно завершает цикл
        result = inconv[result_key]
        break
print(result)
→ Ссылка
Автор решения: Stanislav Volodarskiy

Константу можно улучшить, O-большое останется прежним. Я не уверен, что это поможет, но ничего лучшего не придумалось.

inconv - массив неудобств по всем возможным бригадам. Считается по упорядоченным ростам учеников.

pred(i) проверяет что можно набрать r бригад c неудобством не более i. Возвращает значение как только набрано r бригад.

Двоичный поиск ведётся в диапазоне [min(inconv), max(inconv)].

n, r, c = map(int, str(input()).split())

h = sorted(int(input()) for _ in range(n))
inconv = [h[j] - h[i] for i, j in zip(range(n - c + 1), range(c - 1, n))]


def pred(i):
    k = 0
    j = 0
    while j < len(inconv):
        if inconv[j] <= i:
            k += 1
            if k >= r:
                return True
            j += c
        else:
            j += 1
    return False


low = min(inconv)
high = max(inconv)

# assert not pred(inconv[:low]) and pred(inconv([high:])
while low < high:
    mid = (low + high) // 2
    if pred(mid):
        # assert pred(inconv([mid:])
        high = mid
        # assert pred(inconv([high:])
        # assert not pred(inconv[:low]) and pred(inconv([high:])
    else:
        # assert not pred(inconv([:mid + 1])
        low = mid + 1
        # assert not pred(inconv([:low])
        # assert not pred(inconv[:low]) and pred(inconv([high:])
        
    # assert not pred(inconv[:low]) and pred(inconv([high:])

# assert not pred(inconv[:low]) and pred(inconv([high:])
# assert low == high
# assert not pred(inconv[:low]) and pred(inconv([low:])

print(low)
→ Ссылка
Автор решения: AnnaBazueva

Как вариант:

# Тесты
# n, r, c = map(int, '8 2 3'.split())
# mass = [170, 205, 225, 190, 260, 130, 225, 160]
# n, r, c = map(int, '7 2 3'.split())
# mass = [1, 1, 2, 2, 2, 3, 3]

n, r, c = map(int, input().split())
mass = [int(input()) for _ in range(n)]

mass.sort()

mass = [
    crew[-1] - crew[0] for i in range(len(mass)+1-c) if (crew := mass[i:i+c])
    ]

resu = set()
while r > 0:
    min_mid = min(mass)
    resu.add(min_mid)
    while r > 0 and min_mid in mass:
        mass.remove(min_mid)
        r -= 1
else:
    print(max(resu))

Для дальнейшей оптимизации можно рассмотреть использование collections.Counter,
чтобы избежать многократного удаления элементов и сделать код более эффективным. Это позволит обрабатывать количество вхождений min_mid за O(1) и значительно улучшить общую производительность.

→ Ссылка