Bit Manipulation and Math

Bit manipulation and number theory problems look unfamiliar, and they reward a small set of tricks. Once you know the tricks, many of these questions take a few lines of code and run in constant or logarithmic time. This chapter covers the bitwise operators and the idioms that matter, XOR-based problems, bitmasks for subsets, and the arithmetic and number-theory tools that appear in interviews: gcd, primes, fast exponentiation and modular arithmetic.

1. Bitwise operators

assert 0b1100 & 0b1010 == 0b1000     # AND: 1 only where both bits are 1
assert 0b1100 | 0b1010 == 0b1110     # OR: 1 where either bit is 1
assert 0b1100 ^ 0b1010 == 0b0110     # XOR: 1 where the bits differ
assert ~0 == -1                      # NOT: in Python, ~x == -x - 1
assert 1 << 3 == 8                   # left shift: multiply by 2^3
assert 16 >> 2 == 4                  # right shift: floor divide by 2^2

Python integers have arbitrary precision and no fixed width, which has two consequences: there is no overflow, and negative numbers behave as if they had infinitely many leading 1 bits. To simulate a fixed width, mask: x & 0xFFFFFFFF.

Properties of XOR that solve many problems:

• x ^ x == 0 and x ^ 0 == x.
• XOR is commutative and associative, so the order does not matter.
• Therefore XOR-ing a list in which every value appears twice except one leaves the one.

2. Essential bit idioms

x = 0b10110

assert (x >> 1) & 1 == 1                    # test bit 1 (the second from the right)
assert x | (1 << 0) == 0b10111              # set bit 0
assert x & ~(1 << 1) == 0b10100             # clear bit 1
assert x ^ (1 << 2) == 0b10010              # toggle bit 2

assert x & (x - 1) == 0b10100               # clear the lowest set bit
assert x & -x == 0b00010                    # isolate the lowest set bit

Why x & (x - 1) clears the lowest set bit: subtracting 1 flips the lowest set bit to 0 and all the zeros below it to 1, so ANDing with the original clears exactly that bit.

Is a number a power of two? A power of two has exactly one set bit, so clearing it gives 0.

def is_power_of_two(n):
    return n > 0 and n & (n - 1) == 0

assert is_power_of_two(16) and is_power_of_two(1)
assert not is_power_of_two(18) and not is_power_of_two(0)

Count set bits (Hamming weight). Repeatedly clear the lowest set bit and count how many times. The loop runs once per set bit.

def count_bits(n):
    count = 0
    while n:
        n &= n - 1
        count += 1
    return count

assert count_bits(11) == 3
assert count_bits(128) == 1
assert count_bits(0) == 0
assert bin(2**20 - 1).count("1") == count_bits(2**20 - 1) == 20

Counting bits for every number from 0 to $n$ uses DP: the count for i is the count for i >> 1 plus its last bit.

def counting_bits(n):
    bits = [0] * (n + 1)
    for i in range(1, n + 1):
        bits[i] = bits[i >> 1] + (i & 1)
    return bits

assert counting_bits(5) == [0, 1, 1, 2, 1, 2]

3. XOR problems

Single number. Every element appears twice except one. XOR everything.

def single_number(nums):
    result = 0
    for x in nums:
        result ^= x
    return result

assert single_number([2, 2, 1]) == 1
assert single_number([4, 1, 2, 1, 2]) == 4

Time $O(n)$, space $O(1)$, better than a hash map's $O(n)$ space. Missing number in 0..n: XOR all indexes and values, or use the sum formula.

def missing_number(nums):
    result = len(nums)
    for i, x in enumerate(nums):
        result ^= i ^ x
    return result

assert missing_number([3, 0, 1]) == 2
assert missing_number([0, 1]) == 2

Two numbers appear once, all others twice. XOR everything to get a ^ b. Pick any set bit of that value: a and b differ at that bit. Partition the numbers by that bit and XOR each group.

def single_number_iii(nums):
    xor_all = 0
    for x in nums:
        xor_all ^= x
    low_bit = xor_all & -xor_all                  # a bit where a and b differ
    a = b = 0
    for x in nums:
        if x & low_bit:
            a ^= x
        else:
            b ^= x
    return sorted([a, b])

assert single_number_iii([1, 2, 1, 3, 2, 5]) == [3, 5]

Swap without a temporary (a ^= b; b ^= a; a ^= b) is a curiosity, and not a good production habit. Mention it only if asked.

4. Bitmasks for subsets

A set of up to about 20 elements can be represented as an integer where bit i means "element i is in the set". Then set operations are bit operations, and iterating all subsets is a loop over 0 .. 2^n - 1.

def subsets_bitmask(nums):
    n = len(nums)
    result = []
    for mask in range(1 << n):                       # every subset
        result.append([nums[i] for i in range(n) if mask >> i & 1])
    return result

assert sorted(subsets_bitmask([1, 2, 3])) == sorted([[], [1], [2], [3], [1, 2], [1, 3], [2, 3], [1, 2, 3]])

Enumerate submasks of a mask efficiently:

def submasks(mask):
    sub = mask
    while sub:
        yield sub
        sub = (sub - 1) & mask
    yield 0

assert sorted(submasks(0b101)) == [0, 1, 4, 5]

Bitmasks also give compact state for dynamic programming over subsets (for example the travelling salesman problem with $n\le 20$: dp[mask][last]), and for tracking visited items in a search. They are a typical tool when $n$ is around 15 to 20.

5. Other bit tricks

Reverse bits of a 32-bit integer, number of 1 bits, bitwise AND of a range (the common prefix of the endpoints), sum of two integers without + (XOR gives the sum without carry, AND shifted left gives the carry):

def reverse_bits(n, width=32):
    result = 0
    for _ in range(width):
        result = (result << 1) | (n & 1)
        n >>= 1
    return result

assert reverse_bits(0b00000010100101000001111010011100) == 0b00111001011110000010100101000000

def range_bitwise_and(left, right):
    shift = 0
    while left != right:                   # find the common binary prefix
        left >>= 1
        right >>= 1
        shift += 1
    return left << shift

assert range_bitwise_and(5, 7) == 4
assert range_bitwise_and(0, 0) == 0

def get_sum(a, b):
    mask = 0xFFFFFFFF
    while b & mask:
        a, b = a ^ b, (a & b) << 1         # sum without carry, carry shifted
    return (a & mask) if b > 0 else a if a < 2**31 else ~(a ^ mask)

assert get_sum(1, 2) == 3
assert get_sum(2, 3) == 5

The last one is mostly a curiosity in Python because of unbounded integers, and is a classic in languages with fixed-width integers. The mask handling above is the Python-specific complication.

6. Greatest common divisor and least common multiple

The Euclidean algorithm: gcd(a, b) = gcd(b, a % b), terminating when b is 0. It runs in $O(\log \min (a,b))$.

def gcd(a, b):
    while b:
        a, b = b, a % b
    return a

def lcm(a, b):
    return a // gcd(a, b) * b              # divide first to limit the intermediate size

import math
assert gcd(48, 18) == 6 == math.gcd(48, 18)
assert lcm(4, 6) == 12
assert gcd(7, 0) == 7

Uses: reducing fractions, finding repeating patterns, "can you measure exactly $z$ litres with jugs of $x$ and $y$" (possible when $z$ is a multiple of $\gcd (x,y)$ and $z\le x+y$).

7. Primes

Primality test by trial division up to $\sqrt{n}$:

def is_prime(n):
    if n < 2:
        return False
    i = 2
    while i * i <= n:
        if n % i == 0:
            return False
        i += 1
    return True

assert [p for p in range(20) if is_prime(p)] == [2, 3, 5, 7, 11, 13, 17, 19]

Sieve of Eratosthenes finds all primes up to $n$ in $O(n\log \log n)$: for each prime, mark its multiples as composite, starting at its square.

def count_primes(n):
    """Number of primes strictly less than n."""
    if n < 3:
        return 0
    sieve = [True] * n
    sieve[0] = sieve[1] = False
    for i in range(2, int(n ** 0.5) + 1):
        if sieve[i]:
            for multiple in range(i * i, n, i):
                sieve[multiple] = False
    return sum(sieve)

assert count_primes(10) == 4
assert count_primes(100) == 25
assert count_primes(2) == 0

Prime factorisation by trial division is $O(\sqrt{n})$:

def prime_factors(n):
    factors, d = {}, 2
    while d * d <= n:
        while n % d == 0:
            factors[d] = factors.get(d, 0) + 1
            n //= d
        d += 1
    if n > 1:
        factors[n] = factors.get(n, 0) + 1
    return factors

assert prime_factors(360) == {2: 3, 3: 2, 5: 1}
assert prime_factors(97) == {97: 1}

8. Fast exponentiation and modular arithmetic

Computing $x^n$ by repeated multiplication is $O(n)$. Exponentiation by squaring is $O(\log n)$: $x^n=(x^{n/2}{)}^2$ for even $n$, and $x\cdot x^{n-1}$ for odd $n$.

def power(x, n):
    if n < 0:
        x, n = 1 / x, -n
    result = 1.0
    while n:
        if n & 1:
            result *= x
        x *= x
        n >>= 1
    return result

assert power(2.0, 10) == 1024.0
assert abs(power(2.0, -2) - 0.25) < 1e-12
assert power(2.1, 3) == 2.1 ** 3 or abs(power(2.1, 3) - 9.261) < 1e-9

def mod_pow(base, exp, mod):
    result = 1
    base %= mod
    while exp:
        if exp & 1:
            result = result * base % mod
        base = base * base % mod
        exp >>= 1
    return result

assert mod_pow(2, 10, 1000) == 24 == pow(2, 10, 1000)
assert mod_pow(3, 200, 13) == pow(3, 200, 13)

Modular arithmetic rules for results "modulo $10^9+7$": addition, subtraction and multiplication can be reduced at each step: (a + b) % m, (a * b) % m. Division needs a modular inverse. For a prime modulus $p$, the inverse of $a$ is $a^{p-2}\mathrm{mod}p$ (Fermat's little theorem), and Python provides pow(a, -1, m) in recent versions.

MOD = 10**9 + 7
inv = pow(3, MOD - 2, MOD)
assert 3 * inv % MOD == 1
assert (10 * inv) % MOD == (10 * pow(3, -1, MOD)) % MOD

Combinations and Pascal's triangle: $\left(\genfrac{}{}{0ex}{}{n}{k}\right)$ can be computed with the recurrence $\left(\genfrac{}{}{0ex}{}{n}{k}\right)=\left(\genfrac{}{}{0ex}{}{n-1}{k-1}\right)+\left(\genfrac{}{}{0ex}{}{n-1}{k}\right)$, or multiplicatively, taking care to stay exact:

def n_choose_k(n, k):
    k = min(k, n - k)
    result = 1
    for i in range(1, k + 1):
        result = result * (n - k + i) // i          # stays an integer at every step
    return result

assert n_choose_k(5, 2) == 10 == math.comb(5, 2)
assert n_choose_k(10, 0) == 1
assert n_choose_k(52, 5) == 2598960

9. Integer arithmetic problems

Reverse an integer with overflow handling, palindrome number without converting to a string, count digits, happy number (detect a cycle with a set or fast and slow pointers), excel column title (base-26 with a twist), fizz buzz. They are warm-ups that test edge cases: zero, negatives, overflow, leading zeros.

def reverse_integer(x):
    sign = -1 if x < 0 else 1
    result = int(str(abs(x))[::-1]) * sign
    return result if -2**31 <= result <= 2**31 - 1 else 0

assert reverse_integer(123) == 321
assert reverse_integer(-120) == -21
assert reverse_integer(1534236469) == 0

def is_palindrome_number(x):
    if x < 0 or (x % 10 == 0 and x != 0):
        return False
    reversed_half = 0
    while x > reversed_half:                   # reverse only half the digits
        reversed_half = reversed_half * 10 + x % 10
        x //= 10
    return x == reversed_half or x == reversed_half // 10

assert is_palindrome_number(121) and is_palindrome_number(1221)
assert not is_palindrome_number(-121) and not is_palindrome_number(10)

def is_happy(n):
    seen = set()
    while n != 1 and n not in seen:
        seen.add(n)
        n = sum(int(d) ** 2 for d in str(n))
    return n == 1

assert is_happy(19) is True
assert is_happy(2) is False

def convert_to_title(n):
    out = []
    while n:
        n, r = divmod(n - 1, 26)               # bijective base 26: there is no zero digit
        out.append(chr(ord("A") + r))
    return "".join(reversed(out))

assert convert_to_title(1) == "A"
assert convert_to_title(28) == "AB"
assert convert_to_title(701) == "ZY"

10. Randomised and probabilistic ideas

Reservoir sampling picks a uniformly random item from a stream of unknown length in one pass with $O(1)$ memory: keep the $i$th item with probability $1/i$.

import random

def reservoir_pick(stream):
    chosen = None
    for i, item in enumerate(stream, start=1):
        if random.randint(1, i) == 1:        # probability 1/i
            chosen = item
    return chosen

counts = {x: 0 for x in range(4)}
random.seed(1)
for _ in range(4000):
    counts[reservoir_pick(range(4))] += 1
assert all(800 < c < 1200 for c in counts.values())    # roughly uniform

Fisher-Yates shuffle produces a uniformly random permutation by swapping each position with a random earlier-or-equal position. Know why a naive "swap with any random index" shuffle is biased.

def shuffle(nums):
    for i in range(len(nums) - 1, 0, -1):
        j = random.randint(0, i)
        nums[i], nums[j] = nums[j], nums[i]
    return nums

assert sorted(shuffle([1, 2, 3, 4, 5])) == [1, 2, 3, 4, 5]

11. Common mistakes

• Assuming fixed-width integers in Python. There is no overflow, so reproduce 32-bit behaviour with masks if the problem requires it.
• Operator precedence: x & 1 == 0 parses as x & (1 == 0). Use parentheses: (x & 1) == 0.
• Right shift of negative numbers is an arithmetic shift in Python (it keeps the sign).
• Forgetting to handle zero and negatives in power-of-two, gcd and reverse-integer problems.
• Integer division with negatives: // floors toward negative infinity in Python, unlike truncation in C++ and Java.
• Overflow in the intermediate a * b before taking a modulus in other languages. Reduce first, or use wider types.
• Using floating point for exact integer results, such as int(n ** 0.5) for large $n$, which can be off by one. Use math.isqrt.
• A sieve upper bound off by one (strictly less than $n$ versus at most $n$).

12. Practice set

1. Single number I, II and III, missing number, find the duplicate.
2. Number of 1 bits, counting bits, power of two, power of four, reverse bits.
3. Sum of two integers, bitwise AND of numbers range, maximum XOR of two numbers (a trie on bits).
4. Subsets (bitmask), gray code, shortest path visiting all nodes (bitmask BFS).
5. GCD of strings, water and jug problem, fraction addition.
6. Count primes, ugly numbers, prime factorisation of a factorial.
7. Pow$(x,n)$, super pow, sqrt$(x)$.
8. Reverse integer, palindrome number, happy number, excel sheet column number and title.
9. Random pick with weight, shuffle an array, linked list random node (reservoir sampling).
10. Count of sorted-array-style combinatorics modulo a prime using factorials and modular inverses.