Неравномерный обмен в алгоритме производительности Python

Question

Группа из k мальчиков должна получать мраморные шарики следующим образом: они сидят по кругу, и каждый мальчик берет на 1 мрамор больше, чем мальчик справа от него, и передает сумку с оставшейся частью шариков вмальчик слева от него.

Лидер начинает с того, что достает 1 мрамор для себя.Он передает сумку мальчику слева от него, который затем удаляет 2 шарика для себя и передает сумку слева от него.Этот мальчик затем берет 3 монеты и передает сумку слева от него, и так далее.Этот процесс продолжается до тех пор, пока мешок не опустеет (последний мальчик, который возьмет из сумки, может не получить столько шариков, сколько должен).

Я хочу получить общее количество шариков, которое ЛИДЕР получает в концепроцесса.

Это то, что у меня есть, работает, но слишком медленно:

def countMarbles(n, k):
    c = 0
    leader = 0
    while n>0:
        for i in range(k):
            c+=1
            if i == 0:
                if c<=n:
                    leader += c
                else:
                    leader += n
            n -= c
    return leader

Answer 1

Мрамор, который вы выбрасываете, равен 1, затем 2, затем 3 ... Есть формула для такой суммы, сумма от 1 до x равна (x)(x + 1) / 2.

Теперь вы получилии хотите знать, сколько проходов вы можете сделать.Это означает получение наибольшего значения x, такого, что (x)(x + 1) / 2 меньше или равно n.

. Мы можем получить это, решив 0 = x^2 + x - 2n.Мы могли бы получить десятичный результат там, поэтому мы должны взять минимальное значение положительного ответа на уравнение.

Как только мы нашли правильный x, мы просто знаем, что каждый k проходов мешка, 1 идет клидер.Сначала он получает 1 мрамор, затем он получает k + 1 мрамор, затем 2k + 1 ...

Если было x проходов, ceil x / k достается лидеру.Выбрав первый проход, который всегда равен 1, мы получим l = ceil(x / k) - 1 проходов, у которых коэффициент ак больше 0: ((k + 1) + (2k + 1) + ... + (lk + 1)) = (1 + 2 + 3 + ... + l) * k + l = (l * (l + 1) / 2) * k + l.

С учетом лидера, начавшегося с 1, решение будет (l * (l + 1) / 2) * k + l + 1

Единственная проблема заключается в том, что происходит с оставшимися в сумке оставшимися шариками.В случае, если они должны были пойти к лидеру, мы также должны принять их во внимание.Чтобы это произошло, x должно быть кратным k, что означает, что мы завершили раунд, поэтому следующим должен был стать лидер, но не хватило шариков, чтобы сделать еще один проход.

Вот реализация на python:

import math

def solve (n, k):
    x = math.floor((-1 + math.sqrt(1 + 8*n)) / 2)
    l = math.ceil(x / k) - 1
    sol = (l * (l + 1) / 2) * k + l + 1
    if x % k == 0 :
        sol += n - (x * (x + 1) / 2)
    return int(sol)