Прогрессии и суммы

Арифметические и геометрические прогрессии и быстрые формулы сумм — чтобы не суммировать в цикле то, что считается за O(1).

Арифметическая прогрессия — последовательность с постоянной разностью; геометрическая — с постоянным отношением. У сумм обеих есть замкнутые формулы.

Зачем это в олимпиадах

Суммы вида 1+2+…+n или 1+2+4+…+2k возникают постоянно: при оценке сложности, в подсчётных формулах, при свёртке диапазонов. Считать их циклом за O(n) — расточительство, когда есть формула за O(1). А при больших n (до 1018) цикл вообще невозможен — спасает только формула. Умение мгновенно сворачивать прогрессии и типовые суммы (квадратов, кубов) — признак зрелого олимпиадника. Заодно разберём, как считать геометрическую сумму по модулю, где деление требует обратного элемента.

Арифметическая прогрессия

Прогрессия a, a+d, a+2d, … с первым членом a и разностью d. Сумма первых n членов: складываем первый и последний, умножаем на число членов, делим пополам (трюк юного Гаусса — пары «с краёв» дают одинаковые суммы):

S = (a₁ + aₙ) · n / 2 = (2a + (n−1)d) · n / 2

def arithmetic_sum(a, d, n):
    last = a + (n - 1) * d
    return (a + last) * n // 2

# 1 + 2 + ... + 100
print(arithmetic_sum(1, 1, 100))      # 5050
# проверка циклом
print(sum(range(1, 101)))
# сумма нечётных 1+3+5+...+99 (50 членов)
print(arithmetic_sum(1, 2, 50))       # 2500 = 50^2

Вывод:

5050
5050
2500

Сумма первых n нечётных чисел равна n2 — красивый факт, мгновенно вытекающий из формулы. Целочисленное деление // 2 здесь корректно: произведение (a₁+aₙ)·n всегда чётно (одно из двух соседних по чётности слагаемых).

Геометрическая прогрессия

Прогрессия a, a·r, a·r2, … с отношением r ≠ 1. Сумма первых n членов выводится приёмом «умножь на r и вычти»:

S = a · (rn − 1) / (r − 1)   (при r = 1 сумма равна a·n)

Вывод: пусть S = a + ar + … + arn−1. Тогда rS = ar + ar2 + … + arn. Вычитая, почти всё сокращается: rS − S = arn − a, откуда S = a(rn−1)/(r−1).

def geometric_sum(a, r, n):
    if r == 1:
        return a * n
    return a * (r ** n - 1) // (r - 1)

# 1 + 2 + 4 + ... + 2^9 (10 членов)
print(geometric_sum(1, 2, 10))       # 1023 = 2^10 - 1
print(sum(2 ** i for i in range(10)))  # проверка
# 1 + 3 + 9 + 27 + 81 (5 членов)
print(geometric_sum(1, 3, 5))        # 121

Вывод:

1023
1023
121

Сумма степеней двойки 1+2+…+2n−1 = 2n−1 — это знаменитое «число с n единичными битами». Формула даёт его за O(log n) (из-за возведения rn), а не за O(n) сложений.

Геометрическая сумма по модулю

В задачах сумму берут по простому модулю. Тут возникает деление на r−1 — а делить по модулю можно только умножая на обратный элемент (раздел про модулярную арифметику!). Используем pow(r, n, mod) для степени и pow(r-1, mod-2, mod) для обратного по Ферма.

MOD = 10**9 + 7

def geometric_sum_mod(a, r, n, mod):
    if r % mod == 1:
        return a % mod * (n % mod) % mod
    num = (pow(r, n, mod) - 1) % mod           # r^n - 1
    inv = pow(r - 1, mod - 2, mod)             # обратный к (r-1)
    return a % mod * num % mod * inv % mod

print(geometric_sum_mod(1, 2, 10, MOD))        # 1023
print(geometric_sum_mod(1, 2, 100, MOD))       # (2^100 - 1) mod p
# сверка с прямым суммированием по модулю
print(sum(pow(2, i, MOD) for i in range(100)) % MOD)

Вывод:

1023
976371284
976371284

Деление на r−1 заменили умножением на обратный — и формула дала тот же остаток, что прямое суммирование. Это типичный приём: всюду, где в формуле есть деление, по модулю ставьте обратный элемент.

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

Полезно помнить две классические формулы. Сумма квадратов 12+22+…+n2 = n(n+1)(2n+1)/6. Сумма кубов 13+…+n3 = (n(n+1)/2)2 — квадрат суммы первых n чисел (красивое тождество!).

def sum_squares(n):
    return n * (n + 1) * (2 * n + 1) // 6

def sum_cubes(n):
    s = n * (n + 1) // 2
    return s * s

print(sum_squares(5), sum(i * i for i in range(1, 6)))     # 55
print(sum_cubes(5), sum(i ** 3 for i in range(1, 6)))      # 225
print(sum_cubes(5) == sum(range(1, 6)) ** 2)               # куб = квадрат суммы

Вывод:

55 55
225 225
True

Разбор задачи: сумма целочисленных частей и блочное разбиение

Олимпиадная жемчужина — сумма ∑i=1n ⌊n/i⌋. Наивно это n слагаемых. Но у ⌊n/i⌋ мало различных значений — порядка 2√n, и они идут блоками: на целом диапазоне i частное постоянно. Сгруппировав равные значения, сумму считают за O(√n). Этот приём («деление на блоки», divisor blocks) ускоряет множество теоретико-числовых сумм.

def sum_floor_div(n):
    # сумма floor(n/i) для i = 1..n за O(sqrt(n))
    total = 0
    i = 1
    while i <= n:
        q = n // i
        # наибольший j, при котором floor(n/j) == q
        j = n // q
        total += q * (j - i + 1)     # блок одинаковых значений
        i = j + 1
    return total

print(sum_floor_div(10))
# сверка прямым суммированием
print(sum(10 // i for i in range(1, 11)))
print(sum_floor_div(10**6))         # быстро даже для миллиона

Вывод:

27
27
13970034

Блочное суммирование дало те же 27, что прямой цикл, но обошлось O(√n) итерациями вместо n. Для n = 106 это около двух тысяч шагов вместо миллиона. Сумма ∑⌊n/i⌋ равна суммарному числу делителей всех чисел до n — красивое тождество, а блочный приём делает её вычислимой даже для n до 1012.

Подводные камни

• Случай r = 1. Формула геометрической суммы делит на r−1 — при r=1 это деление на ноль; обрабатывайте отдельно (S = a·n).
• Деление по модулю. /(r−1) по модулю — это умножение на обратный, а не //.
• Чётность в арифметической сумме. (a₁+aₙ)·n всегда чётно, поэтому //2 точно; не теряйте этого при модуле (там тоже нужен обратный к 2).
• Большие n. rn считайте через pow, а не r ** n в цикле — иначе медленно и без модуля огромно.

Итог

• Арифметическая сумма: (a₁+aₙ)·n/2; сумма первых n нечётных = n2.
• Геометрическая сумма: a(rn−1)/(r−1), отдельно r=1.
• По модулю деление на r−1 заменяется обратным элементом по Ферма.
• Суммы квадратов n(n+1)(2n+1)/6 и кубов (n(n+1)/2)2 — наизусть.