1 minute read

Increment an integer as array - Elements of Programming Interviews (EPI)

Problem:

Write a program which takes an array of integers that represent an integer I and updates the array to I+1. Example: [9,9,9,9] should return [1,0,0,0,0], [1,2,3,4] should return [1,2,3,5].

Solution:

The brute force or naive approach would be to first convert the array into an integer, add 1 and then convert the result to an array again. This is limited by the size of the integer type. A better solution would be to simply add 1 to the last digit and carry forward if it is equal to 10. We used a very hacky approach with a try and except block:

def plus_one(A: List[int]) -> List[int]:
    """Function to increment an integer
    represented as an array."""
    carry = 1
    i = -1
    try:
        while carry:
            # Traverse the array in reverse and
            # keep adding the digits.
            carry = (A[i] + 1) // 10
            A[i] = (A[i] + 1) % 10
            i -= 1
    except IndexError:
        # We have crossed the entire length of
        # the array like 999 +1 -> 000 so we
        # simply append a 0 at the end and set
        # first element to 1 -> 1000
        # This can only happen if the entire array
        # is 9s. (9999...)
        A.append(0)
        A[0] = 1
    return A

A much cleaner solution from the book is as follows:

def plus_one(A: List[int]) -> List[int]:

    A[-1] += 1
    for i in reversed(range(1, len(A))):
        if A[i] != 10:
            break
        A[i] = 0
        A[i - 1] += 1
    else:
        if A[0] == 10:
            # There is a carry-out, so we need one more digit to store the result.
            # A slick way to do this is to append a 0 at the end of the array,
            # and update the first entry to 1.
            A[0] = 1
            A.append(0)
    return A