0

Pascal triangle digit encryption

by · 1 day ago· ⑂ 0 remixes

Question: Pascal triangle digit encryption

Companies: Amazon

Difficulty: Easy

Problem statement

Digits are encrypted by folding them down an upside-down Pascal's triangle.

You get a list of n digits (each 0-9). In one round, replace the list with the sums of each pair of neighbours, keeping only the last digit of each sum ((a + b) % 10). The list gets one shorter each round. Stop when exactly two digits remain.

Implement:

def encrypt(digits: list[int]) -> str

Return the final two digits as a two-character string. Keep a leading zero ("07", not "7").

Constraints

  • 2 <= n <= 1000 (a follow-up raises this to 10^5)
  • every element is between 0 and 9

Example 1

digits = [1, 2, 3, 4]
round 1: [3, 5, 7]
round 2: [8, 2]
output: "82"

Example 2

digits = [0, 7]
already two digits
output: "07"

Approach

Each round replaces the row with the sums of neighbouring digits, keeping only
the last digit ((a + b) % 10). The row gets one shorter per round, and we
stop at two digits. The answer is those two digits as a string, so a leading
zero stays ("07", not "7").

Simulation (what most interviewers expect first). Keep a list and rebuild
it n - 2 times. Each round is linear, so this costs O(n^2) time. You can
overwrite the list in place from left to right (row[i] = (row[i] + row[i+1]) % 10),
because row[i+1] is still the old value when you read it. That gives O(1)
extra space. For typical online-assessment limits (n up to a few thousand)
this is the whole answer.

Closed form (the follow-up for large n). The process is Pascal's triangle
upside down. After m rounds, position j holds
sum_{i=0..m} C(m, i) * a[j + i] (mod 10), because each digit reaches
position j along C(m, i) paths. With m = n - 2 rounds:

  • first digit = sum C(n-2, i) * a[i] mod 10
  • second digit = sum C(n-2, i) * a[i + 1] mod 10

We cannot divide mod 10 (10 is not prime), so compute C(m, i) mod 10 from
its residues mod 2 and mod 5:

  • mod 2: by Lucas's theorem, C(m, k) is odd exactly when k & m == k.
  • mod 5: Lucas's theorem again, digit by digit in base 5, with a 5x5 table
    of small binomials.
  • combine with the Chinese remainder theorem: the unique x in 0..9 with
    x % 2 == r2 and x % 5 == r5 is r5 or r5 + 5.

Each coefficient costs O(log_5 m), so the whole thing is O(n log n).

Invariant (simulation): after round r, the row has n - r digits and
row[j] equals the digit at position j of level r of the triangle.

Code

from typing import List, Sequence


def encrypt_simulate(digits: Sequence[int]) -> str:
    """O(n^2) time, O(n) space for the copy (O(1) extra if done in place)."""
    if len(digits) < 2:
        raise ValueError("need at least two digits")
    row: List[int] = list(digits)
    length = len(row)
    while length > 2:
        for i in range(length - 1):
            row[i] = (row[i] + row[i + 1]) % 10  # row[i+1] is still old
        length -= 1
    return f"{row[0]}{row[1]}"


# Small binomials C(a, b) for 0 <= a, b < 5, used by Lucas mod 5.
_SMALL = [[0] * 5 for _ in range(5)]
for _a in range(5):
    _SMALL[_a][0] = 1
    for _b in range(1, _a + 1):
        _SMALL[_a][_b] = _SMALL[_a - 1][_b - 1] + _SMALL[_a - 1][_b]


def _binom_mod2(m: int, k: int) -> int:
    return 1 if (k & m) == k else 0


def _binom_mod5(m: int, k: int) -> int:
    result = 1
    while m or k:
        mi, ki = m % 5, k % 5
        if ki > mi:
            return 0
        result = (result * _SMALL[mi][ki]) % 5
        m //= 5
        k //= 5
    return result


def _binom_mod10(m: int, k: int) -> int:
    r2, r5 = _binom_mod2(m, k), _binom_mod5(m, k)
    return r5 if r5 % 2 == r2 else r5 + 5  # CRT for moduli 2 and 5


def encrypt_fast(digits: Sequence[int]) -> str:
    """O(n log n) time, O(1) extra space."""
    n = len(digits)
    if n < 2:
        raise ValueError("need at least two digits")
    m = n - 2
    first = second = 0
    for i in range(m + 1):
        c = _binom_mod10(m, i)
        if c:
            first = (first + c * digits[i]) % 10
            second = (second + c * digits[i + 1]) % 10
    return f"{first}{second}"


if __name__ == "__main__":
    import random

    assert encrypt_simulate([1, 2, 3, 4]) == "82"
    assert encrypt_simulate([9, 9, 9]) == "88"
    assert encrypt_simulate([0, 7]) == "07"
    assert encrypt_simulate([5, 5, 5, 5, 5]) == "00"
    for _ in range(2000):
        arr = [random.randint(0, 9) for _ in range(random.randint(2, 60))]
        assert encrypt_simulate(arr) == encrypt_fast(arr), arr
    print("ok")

Complexity

  • encrypt_simulate: O(n^2) time (n - 2 rounds of linear work).
    O(n) space for the working copy; O(1) extra if you may overwrite the
    input.
  • encrypt_fast: O(n log n) time (n - 1 coefficients, each
    O(log_5 n) for Lucas mod 5 and O(1) mod 2). O(1) extra space.

Test cases

  • [1, 2, 3, 4] -> [3, 5, 7] -> [8, 2] -> "82"
  • [9, 9, 9] -> [8, 8] -> "88"
  • [0, 7] (already two digits) -> "07" (leading zero kept)
  • [5, 5, 5, 5, 5] -> [0, 0, 0, 0] -> ... -> "00"
  • [1, 0, 0, 0, 0, 0] -> first digit is C(4, 0) = 1, second is 0 -> "10"
  • [0, 0, 0, 0, 0, 1] -> first 0, second C(4, 4) = 1 -> "01"
  • [1, 1, 1, 1, 1, 1, 1] (m = 5) -> each digit is 2^5 mod 10 = 2 -> "22"
  • Fewer than two digits -> error (or whatever the interviewer specifies)
  • Random arrays up to length 60: simulation and closed form agree

Follow-ups

  1. Input as a string, return the same format. Parse with
    int(ch) per character; nothing else changes. Confirm the output keeps
    leading zeros.
  2. Return only whether the two final digits are equal (a common variant).
    Compare first == second at the end; same complexity.
  3. n up to 10^5 or 10^6. The quadratic loop times out; switch to the
    binomial form with Lucas mod 2 and mod 5 plus CRT. Explain why plain
    modular inverses fail: 10 is not prime and C(m, k) has factors of 2 and 5.
  4. Stop at one digit, or at k digits, or use a different base B. Run
    n - k rounds, so use m = n - k and compute k sums. For base B,
    factor B into prime powers; prime moduli use Lucas, prime powers need
    Granville's extension or a factor-counting trick (count the powers of p in
    C(m, k) incrementally with Kummer/Legendre and keep the unit part).
  5. Many queries on subarrays. Each query [l, r] uses m = r - l - 1,
    so the coefficients depend only on the length. Precompute coefficient
    rows per distinct length, or answer each query in O(len log len).

Interviewer expectations

Mid-level (L4)

  • Simulates correctly, rounds stop at exactly two digits, keeps the leading zero.
  • Handles the two-digit input and states the O(n^2) cost.
  • Writes a few hand-checked tests.

Senior (L5)

  • Does the update in place and explains why reading row[i+1] is safe.
  • Spots that the result is a binomial-weighted sum (Pascal's triangle) and
    says why that matters for large n.
  • Handles the "are the digits equal" and string-input variants quickly.

Staff (L6+)

  • Derives the O(n log n) method: Lucas's theorem mod 2 and mod 5, CRT to
    mod 10, and explains why division mod 10 is not an option.
  • Checks the fast version against the simulation with a randomized test.
  • Generalises to other bases or final lengths and discusses trade-offs
    (simple loop for OA limits vs number theory when n is large).

Sign in to comment, vote or remix.