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 to10^5) - every element is between
0and9
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 whenk & 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
xin0..9with
x % 2 == r2andx % 5 == r5isr5orr5 + 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 - 2rounds 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 - 1coefficients, each
O(log_5 n)for Lucas mod 5 andO(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 isC(4, 0) = 1, second is0->"10" -
[0, 0, 0, 0, 0, 1]-> first0, secondC(4, 4) = 1->"01" -
[1, 1, 1, 1, 1, 1, 1](m = 5) -> each digit is2^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
-
Input as a string, return the same format. Parse with
int(ch)per character; nothing else changes. Confirm the output keeps
leading zeros. -
Return only whether the two final digits are equal (a common variant).
Comparefirst == secondat the end; same complexity. -
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 andC(m, k)has factors of 2 and 5. -
Stop at one digit, or at k digits, or use a different base B. Run
n - krounds, so usem = n - kand computeksums. 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). -
Many queries on subarrays. Each query
[l, r]usesm = r - l - 1,
so the coefficients depend only on the length. Precompute coefficient
rows per distinct length, or answer each query inO(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).