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

Есть код

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


def calculate(m, grades):
    result = 0

    for i in range(m):
        if grades[i] != 0:
            result += grades[i]**2
            counter = 0
            for j in range(i + 1, m):
                if counter == grades[i]:
                    break
                if grades[j] != 0:
                    result += grades[j]
                    counter += 1
        return result


print(calculate(m, grades))

Как убрать в нем циклы, без потери функционала кода, так сказать. Условие задачи