Решение задачи с Leetcode про возведение числа в степень / Хабр

Данная статья является продолжением предыдущей статьи и также посвящена использованию битовой арифметики для решения задач с LeetCode.

На этот раз рассмотрим задачу 50. Pow(x, n) в которой нас просят реализовать алгоритм возведения произвольного числа x в целую степень n. Ограничения заданные в условиях:

$-100.0 \le x \le 100.0 \\ -2^{31} \le n \le 2^{31}-1$

Очевидно, что первым на ум приходит наивный алгоритм, где нам надо сделать n-1 умножение и вернуть результат. Этот алгоритм работает за O(n) и прост в понимании.

Но, если нам предстоит много операций возведения в степень, то хочется его как-то улучшить. Для этого можно воспользоваться тем фактом, что любое положительное n мы можем представить в бинарном виде, для того, чтобы урезать вычислительную сложность до O(log n)

$a^n = a^{b_0 \cdot 2^0 + b_1 \cdot 2^1 + \dots + b_k \cdot 2^k} = \prod_{i=0}^k a^{b_i \cdot 2^i}$

Осталось реализовать это решение:

def my_pow(x: float, n: int) -> float:
    if n == 0:
        return 1.0

    # Обработка отрицательной степени
    if n < 0:
        x = 1 / x
        n = -n

    result = 1.0
    current_product = x

    while n > 0:
        if n & 1:
          # нашли отличный от 1 множитель в нашем разложении
          result *= current_product
        current_product *= current_product
        
        n = n >> 1
    
    return result

В результате, вместо n-1 умножений наивного алгоритма, нам надо сделать всего лишь количество итераций, равное бинарному представлению числа n, что дает вычислительную сложность O(log n) и константную сложность по памяти.

Теперь, когда мы умеем быстро возводить в произвольную степень, давайте решим похожую задачу: 372. Super Pow в которой нас просят посчитать

$a^{b} mod \space 1337$

при условии, что b - экстремально большое положительное целое число, заданное в виде массива, а также:

$1 \le a \le 2^{31} - 1 \\ 1 \le b.lenght \le 2000 \\ 0 \le b[i] \le 9 \\ b \space не \spaceсодержит \spaceведущих \spaceнулей$

Основная сложность заключается в том, что b - очень большое, но нам может помочь тот факт, что в задаче от нас требуется предоставить результат деления по модулю в качестве ответа, что означает, что ответ будет в диапазоне от 0 до 1337

Опять воспользуемся тем фактом, что мы способны представить b в виде произведения степеней 2.

class Solution:
    def super_pow(self, a: int, b: list[int]) -> int:
        MOD = 1337

        def pow_mod(base: int, exponent: int) -> int:
            result = 1
            base = base % MOD
            while exponent > 0:
                if exponent & 1:
                    result = (result * base) % MOD
                base = (base * base) % MOD
                exponent = exponent >> 1
            return result

        result = 1
        a = a % MOD  

        for digit in b:
            result = (pow_mod(result, 10) * pow_mod(a, digit)) % MOD

        return result

В результате получаем алгоритм, который за O(n log k), где n - число десятичных разрядов числа b (длина массива), а k - количество разрядов в бинарном представлении цифры, то есть от 0 до 4. Отметим также, что этот алгоритм использует константную дополнительную память.