Printing Pascal’s triangle for n number of rows in Python

Question

I have written a program that takes a number n and prints Pascal's triangle having n number of rows.

Here is my code -

n = int(input("Enter number of rows: "))
a = []

for i in range(n):
    a.append([])
    a[i].append(1)

    for j in range(1, i):
        a[i].append(a[i - 1][j - 1] + a[i - 1][j])
    if (n != 0):
        a[i].append(1)

for i in range(n):
    print("   " * (n - i), end = " ", sep = " ")

    for j in range(0, i + 1):
        print('{0:5}'.format(a[i][j]), end = " ", sep = " ")
    print()

Here are some example outputs -

Enter number of rows: 3
              1 
           1     1 
        1     2     1 

---

Enter number of rows: 6
                       1 
                    1     1 
                 1     2     1 
              1     3     3     1 
           1     4     6     4     1 
        1     5    10    10     5     1 

---

Enter number of rows: 10
                                   1 
                                1     1 
                             1     2     1 
                          1     3     3     1 
                       1     4     6     4     1 
                    1     5    10    10     5     1 
                 1     6    15    20    15     6     1 
              1     7    21    35    35    21     7     1 
           1     8    28    56    70    56    28     8     1 
        1     9    36    84   126   126    84    36     9     1 

Therefore, I would like to know whether I could make this program shorter and more efficient.

Any help would be highly appreciated.

Answer 1

One issue with your approach is your memory usage is \$O(n^2)\$. You are building up the entire Pascal's Triangle, and subsequently printing it. If you simply want to print Pascal's Triangle, you don't need to keep the entire triangle in memory. You can simply generate and print it row by row. Since you generate the next row from only the previous row, the memory requirement is simply two rows, which is \$O(n)\$.

>>> row = [1]
>>> for _ in range(5):
...    row = [1] + [x + y for x, y in zip(row[:-1], row[1:])] + [1]
...    print(row)
... 
[1, 1]
[1, 2, 1]
[1, 3, 3, 1]
[1, 4, 6, 4, 1]
[1, 5, 10, 10, 5, 1]
>>> 

Demystifying [1] + [x + y for x, y in zip(row[:-1], row[1:])] + [1]:

row[:-1] is all the values in row, except the last one. For example, if row is [1, 3, 3, 1], then row[:-1] is [1, 3, 3]. Similarly, row[1:] is all the values of row except the first, so [3, 3, 1]. zip() takes the first value of each of the lists, and returns those as the tuple (1, 3), which get assigned to x and y respectively, which get added together to get 4. Then, zip() emits the next pair of values, (3, 3) which results in 6, and then the final pair (3, 1) results in 4, producing the list [4, 6, 4]. The [1] + ... + [1] adds the bookends, which results in the next row of the triangle: [1, 4, 6, 4, 1].

Instead of needing to keep track of the last row, and feed that back to generate the next row, we can package that up into a generator function, which can prime the loop and do the feedback itself:

>>> def pascals_triangle(n):
...     row = [1]
...     yield row
...     for _ in range(n):
...         row = [1] + [x + y for x, y in zip(row[:-1], row[1:])] + [1]
...         yield row
... 
>>> for row in pascals_triangle(5):
...     print(row)
... 
[1]
[1, 1]
[1, 2, 1]
[1, 3, 3, 1]
[1, 4, 6, 4, 1]
[1, 5, 10, 10, 5, 1]
>>> 

Format the resulting rows as desired.

>>> for i, row in enumerate(pascals_triangle(5)):
...     print(("   "*(5-i) + "{:6}"*(i+1)).format(*row))
... 
                    1
                 1     1
              1     2     1
           1     3     3     1
        1     4     6     4     1
     1     5    10    10     5     1
>>> 

Answer 2

Functions & code organisation

At the moment, your code is basically split into 3 parts:

• getting the input from the user
• generating a triangle
• printing a triangle

Things would be easier to understand/test/reuse if we could split them into 3 logical parts. Functions could be a good way to do things here.

Also, it could be a good occasion to use the if __name__ == "__main__": trick to split function definitions from the code using it and performing Input/Output.

def generate_pascal(n):
    ...
    return a

def print_pascal(a):
    ...

if __name__ == "__main__":
    n = int(input("Enter number of rows: "))
    p = generate_pascal(n)
    print_pascal(p)

---

Simplifying print_pascal

You can define end and sep as parameters to print_pascal.

Loop like a native: I highly recommend Ned Batchelder's excellent talk called "Loop like a native". You usually do not need to write loops based on the length of the list you are working on, you can just iterate over it.

In our case, this leads to the following code:

def print_pascal(a, end = " ", sep = " "):
    n = len(a)
    for i, line in enumerate(a):
        print("   " * (n - i), end = end, sep = sep)
        for e in line:
            print('{0:5}'.format(e), end = sep, sep = sep)
        print()

Note: This hilight an issue with the structure provided to the printing procedure as the first line contains 2 numbers.

---

Simplifying generate_pascal

In generate_pascal, you add a list to a there refer to that list using a[i].

It would be clearer and more efficient to just define a new list that you fill and then eventually add to a.

We get:

def generate_pascal(n):
    a = []
    for i in range(n):
        line = []
        line.append(1)
        for j in range(1, i):
            line.append(a[i - 1][j - 1] + a[i - 1][j])
        if (n != 0):
            line.append(1)
        a.append(line)
    return a

Fixing a bug in generate_pascal

The issue hilighted previously can probably be fixed by replacing if (n != 0) with if (i != 0) which is more usually written without the superfluous parenthesis: if i != 0 or even if i.

More generally, this if could be applied to most of the logic in the loop:

def generate_pascal(n):
    a = []
    for i in range(n):
        line = []
        line.append(1)
        if i:
            for j in range(1, i):
                line.append(a[i - 1][j - 1] + a[i - 1][j])
            line.append(1)
        a.append(line)
    return a

Now, we can take this chance to access a[i - 1] via a variable with a proper name for instance prev_line. Also, The conventional way to access the latest element from an array is to use the -1 index. We could write:

        prev_line = a[-1]
        for j in range(1, i):
            line.append(prev_line[j - 1] + prev_line[j])

---

This this stage, the full code looks like:

def generate_pascal(n):
    p = []
    for i in range(n):
        line = [1]
        if i:
            prev_line = p[-1]
            for j in range(1, i):
                line.append(prev_line[j - 1] + prev_line[j])
            line.append(1)
        p.append(line)
    return p

def print_pascal(p, end = " ", sep = " "):
    n = len(p)
    for i, line in enumerate(p):
        print("   " * (n - i), end = end, sep = sep)
        for e in line:
            print('{0:5}'.format(e), end = sep, sep = sep)
        print()

if __name__ == "__main__":
    n = int(input("Enter number of rows: "))
    p = generate_pascal(n)
    print_pascal(p)

Answer 3

There's an even better way to generate the Pascal's triangle that only requires O(1) memory. It can even calculate a certain row directly without requiring previous rows to be computed.

It works like this:

• the first number in each row is always 1.
• consider the fraction (n / 1) where n in the row index (1-based)
• the 2nd number is obtained by multiplying the first number by this fraction
• decrement the numerator of the fraction by 1 and increment the denominator by 1
• obtain the 3rd number by multiplying the fraction with the 2nd number
• and so on, for (n + 1) numbers

Worked Example:

suppose n = 4

• 1st number = 1
• 2nd number = 1 * (4 / 1) = 4
• 3rd number = 4 * (3 / 2) = 6
• 4th number = 6 * (2 / 3) = 4
• 5th number = 4 * (1 / 4) = 1

which is exactly the 4th row of triangle

Code:

def print_row(n):
    numerator, denominator = n, 1
    cur = 1
    for _ in range(n + 1):
        print(cur, end='\t')
        cur = (cur * numerator) // denominator
        numerator -= 1
        denominator += 1
    print()

for row in range(5):
    print_row(row)

Answer 4

The following is a profiling run of all of the methods I've seen above, with some improvements to the I/O.

from io import StringIO
from timeit import timeit
from sys import stdout, stderr
import numpy as np


def justin(n):
    a = []

    for i in range(n):
        a.append([])
        a[i].append(1)

        for j in range(1, i):
            a[i].append(a[i - 1][j - 1] + a[i - 1][j])
        if n != 0:
            a[i].append(1)

    for i in range(n):
        print("   " * (n - i), end=" ", sep=" ")

        for j in range(0, i + 1):
            print('{0:5}'.format(a[i][j]), end=" ", sep=" ")
        print()


def josay(n):
    def generate_pascal():
        p = []
        for i in range(n):
            line = [1]
            if i:
                prev_line = p[-1]
                for j in range(1, i):
                    line.append(prev_line[j - 1] + prev_line[j])
                line.append(1)
            p.append(line)
        return p

    def print_pascal(p, end=" ", sep=" "):
        for i, line in enumerate(p):
            print("   " * (n - i), end=end, sep=sep)
            for e in line:
                print('{0:5}'.format(e), end=sep, sep=sep)
            print()

    print_pascal(generate_pascal())


def neufeld_pascals_triangle(n):
    row = [1]
    yield row
    for _ in range(n - 1):
        row = [1] + [x + y for x, y in zip(row[:-1], row[1:])] + [1]
        yield row


def neufeld(n):
    for i, r in enumerate(neufeld_pascals_triangle(n)):
        print(("   "*(n-i) + "{:6}"*(i+1)).format(*r))


def neufeld_stdout(n):
    for i, r in enumerate(neufeld_pascals_triangle(n)):
        stdout.write(('   '*(n-i-1) + '{:<6}'*(i+1) + '\n').format(*r))


def neufeld_noenum(n):
    for r in neufeld_pascals_triangle(n):
        i = len(r)
        stdout.write(('   '*(n-i) + '{:<6}'*i + '\n').format(*r))


def neufeld_fmt(n):
    fmt = '   '*n
    for r in neufeld_pascals_triangle(n):
        fmt = fmt[3:] + '{:<6}'
        stdout.write((fmt + '\n').format(*r))


def neufeld_onewrite(n):
    fmt = '   '*n
    msg = ''
    for r in neufeld_pascals_triangle(n):
        fmt = fmt[3:] + '{:<6}'
        msg += fmt.format(*r) + '\n'
    stdout.write(msg)


def neufeld_strio(n):
    fmt = '   '*n
    msg = StringIO()
    for r in neufeld_pascals_triangle(n):
        fmt = fmt[3:] + '{:<6}'
        msg.write(fmt.format(*r) + '\n')
    stdout.write(msg.getvalue())


def tri_rishav(n):
    for y in range(n):
        numerator, denominator = y, 1
        cur = 1
        row = [None] * (y + 1)
        for x in range(y + 1):
            row[x] = cur
            cur = (cur * numerator) // denominator
            numerator -= 1
            denominator += 1
        yield row
def rishav_strio(n):
    fmt = '   '*n
    msg = StringIO()
    for r in tri_rishav(n):
        fmt = fmt[3:] + '{:<6}'
        msg.write(fmt.format(*r) + '\n')
    stdout.write(msg.getvalue())


def tri_numpy(n):
    row = []
    for _ in range(n):
        new_row = np.ones(1 + len(row), np.int32)
        new_row[1:-1] = row[1:] + row[:-1]
        row = new_row
        yield row
def numpy_strio(n):
    fmt = '   '*n
    msg = StringIO()
    for r in tri_numpy(n):
        fmt = fmt[3:] + '{:<6}'
        msg.write(fmt.format(*r) + '\n')
    stdout.write(msg.getvalue())


def tri_onearray(n):
    row = np.ones(n, np.int32)
    x = n - 1
    for y in range(n):
        yield row[x: x+y+1]
        row[x: x+y] += row[x+1: x+y+1]
        x -= 1
def onearray_strio(n):
    fmt = '   '*n
    msg = StringIO()
    for r in tri_onearray(n):
        fmt = fmt[3:] + '{:<6}'
        msg.write(fmt.format(*r) + '\n')
    stdout.write(msg.getvalue())


methods = ((n, globals()[n]) for n in (
    'justin', 'josay', 'neufeld',
    'neufeld_stdout', 'neufeld_noenum', 'neufeld_fmt', 'neufeld_onewrite',
    'neufeld_strio', 'rishav_strio',
    'numpy_strio', 'onearray_strio'))


def profile():
    # IMPORTANT - run redirecting stdout to /dev/null !
    first = None
    for name, method in methods:
        method = globals()[name]
        n = 50

        def run():
            method(n)
        reps = 200
        stdout.flush()
        dur = timeit(run, number=reps) / reps
        if first is None:
            first = dur
            rel = 0
        else:
            rel = first/dur - 1
        print(f'n={n} method={name:16} {1e3*dur:5.3f}ms speedup={rel:7.2%}',
              file=stderr, flush=True)


def print_test():
    for name, method in methods:
        method(6)
        print(name)


profile()
# print_test()

This yields the following results:

n=50 method=justin           7.276ms speedup=  0.00%
n=50 method=josay            6.956ms speedup=  4.60%
n=50 method=neufeld          2.110ms speedup=244.85%
n=50 method=neufeld_stdout   1.947ms speedup=273.78%
n=50 method=neufeld_noenum   1.951ms speedup=273.02%
n=50 method=neufeld_fmt      1.925ms speedup=277.90%
n=50 method=neufeld_onewrite 1.577ms speedup=361.45%
n=50 method=neufeld_strio    1.555ms speedup=367.87%
n=50 method=rishav_strio     1.714ms speedup=324.50%
n=50 method=numpy_strio      2.687ms speedup=170.80%
n=50 method=onearray_strio   2.529ms speedup=187.66%

Of note:

• It was quite interesting to me that neither Rishav's computed method nor Numpy vectorization are able to beat Neufeld's simple, array-based method
• Building up the printed output in memory before writing it to stdout is much faster than writing to stdout on each iteration, especially if you avoid print