Vertically print a list of strings to STDOUT into 3 columns with column lengths as balanced as possible

Question

This was a interview question I struggled with, I am posting my solution hoping to look for more elegant/efficient solutions.

Problem: Given a list of strings, print its element out to STDOUT one by one vertically into 3 columns, while making sure the 3 columns' length are as balanced as possible.

For example, given list = ['a','b','c','d','e','f','g','h','i','j']

The output could be either:

a e h
b f i
c g j
d

OR

a d g
b e h
c f i
    j

But not:

a e i
b f j
c g 
d h  

My approach in Python 3:

Basically I put each word into a "N by 3" 2D list horizontally and try to print them out vertically. If the length is not divisible by 3, I manually increase the length of the first 2 rows depending on the remainder. This creates quite a few edge cases for me to handle separately, such as when the length of the list is between 3 and 6 or when the length is less than or equal to 3.

def assign_words(grid, words):
    i = 0
    for row in range(len(grid)):
        for col in range(len(grid[row])):
            grid[row][col] = words[i]
            i += 1
            if i == len(words):
                return

def print_words(words):
    word_count = len(words)
    rows = word_count // 3
    extra_columns = word_count % 3

    grid = [[''] * 3 for _ in range(rows)]

    # special case
    if word_count <= 3:
        grid.append(['']*3)
        assign_words(grid, words)
        for row in grid:
            print (' '.join(row))
        return

    if extra_columns == 1:
        if rows > 2:
            grid[0].append('')
        elif rows == 2:
            grid[1].pop()
            grid.append(['']*3)
    elif extra_columns == 2:
        if rows > 2:
            grid[0].append('')
            grid[1].append('')
        else:
            grid.append(['']*3)

    assign_words(grid, words)

    # special case
    if 3 < word_count < 6:
        print (grid[0][0]+ ' '+grid[0][2]+' '+grid[1][1])
        print (grid[0][1]+ ' '+grid[1][0]+' '+grid[1][2])
        return

    # print grid
    for col in range(len(grid[0])):
        for row in range(len(grid)):
            if col == len(grid[row]):
                break
            print (grid[row][col], end=' ')
        print ()
 print_words(['a','b','c','d','e','f','g'])

Answer 1

You should take a look into Python advanced iteration capabilities, i.e. the itertools module. You’ll find a number of recipe; amongst which is take:

import itertools


def take(n, iterable):
    "Return first n items of the iterable as a list"
    return list(itertools.islice(iterable, n))

Which will help you extract columns out of your list of words. You can divide the length of the list by 3 to know how much data to take in each column, and adjust using the modulo of the length by 3 (the modulo first columns get to pick one more word); which can be eased thanks to divmod. Lastly, you need to organize your columns into rows, which is easily done using itertools.zip_longest.

That being done, you will need to print the data. The simplest solution is to only print them, as you do:

def print_words(words):
    columns, dangling = divmod(len(words), 3)
    iterator = iter(words)
    columns = [take(columns + (dangling > i), iterator) for i in range(3)]
    for row in itertools.zip_longest(*columns, fillvalue=''):
        print('{}  {}  {}'.format(*row))

Which behaves like so:

>>> print_words(['a', 'b', 'ccccccc', 'dd', 'e', 'ffffffff', 'gg', 'h', 'i', 'j'])
a  e  h
b  ffffffff  i
cccccc  gg  j
dd

---

Which is not pretty. So you might want to save the intermediate result to check, for each column, the length of the longest word, and apply some padding:

def print_words(words):
    columns, dangling = divmod(len(words), 3)
    iterator = iter(words)
    columns = [take(columns + (dangling > i), iterator) for i in range(3)]
    paddings = [max(map(len, column)) for column in columns]
    for row in itertools.zip_longest(*columns, fillvalue=''):
        print('  '.join(word.rjust(pad) for word, pad in zip(row, paddings)))

Which behave more nicely:

>>> print_words(['a', 'b', 'ccccccc', 'dd', 'e', 'ffffffff', 'gg', 'h', 'i', 'j'])
      a         e  h
      b  ffffffff  i
ccccccc        gg  j
     dd             

---

Lastly, we can generalize to support an arbitrary number of columns:

def print_words(words, column_count=3):
    columns, dangling = divmod(len(words), column_count)
    iterator = iter(words)
    columns = [take(columns + (dangling > i), iterator) for i in range(column_count)]
    paddings = [max(map(len, column)) for column in columns]
    for row in itertools.zip_longest(*columns, fillvalue=''):
        print('  '.join(word.rjust(pad) for word, pad in zip(row, paddings)))

Which allows

>>> print_words(['a', 'b', 'ccccccc', 'dd', 'e', 'ffffffff', 'gg', 'h', 'i', 'j'])
      a         e  h
      b  ffffffff  i
ccccccc        gg  j
     dd             
>>> print_words(['a', 'b', 'ccccccc', 'dd', 'e', 'ffffffff', 'gg', 'h', 'i', 'j'], 4)
      a        dd  gg  i
      b         e   h  j
ccccccc  ffffffff       
>>> print_words(['a', 'b', 'ccccccc', 'dd', 'e', 'ffffffff', 'gg', 'h', 'i', 'j'], 5)
a  ccccccc         e  gg  i
b       dd  ffffffff   h  j

Answer 2

There's two separate things to work through here, clarity of code and using appropriate algorithm. I will treat them separately.

algorithm

You offer a somewhat open-ended spec, observing that we could print This or That with both being equally correct. That's fine. But you'll need to nail down details before settling on a concrete algorithm. So arbitrarily pick one. I recommend the first output option, as it allows for iterating in the most natural way.

You suggest that you need a 2-D data structure, offering random access, in order to supply correct output. I don't see why that is necessary. Stick with simple if you can.

option 1

Create an array of 3 'current' indices pointing initially at 'a', 'e', & 'h'. Create an array of 3 'end' indexes pointing at 'e', 'h', and past end of array.

For each row of output, for each column of output, if current less than end, emit current word and increment index.

option 2

At the cost of some extra storage we can simplify further.

Compute stride as approximately one-third the length of words, and initialize done = set().

For each row of output i, step across each column of output by the stride, so you conditionally output word[j]. Add j to done. Suppress output when j in done.

option 3

Don't use any storage. There are three cases, according to whether number of words mod 3 is 0, 1, or 2. Just hardcode it with ifs, or use div & mod expressions.

The first two options are focused on making it easy to test that row 0 is output correctly, and that transitioning to each following row is clearly handled correctly. The third option is best, at the expense of ease of testing.

code

def assign_words(grid, words):

This might be a very nice function. Hard to tell, since you didn't offer a comment or docstring describing what a correct assignment of words would look like. There's two ways to exit, through the conditional or through completing the nested loops, so there is a special relationship between input array dimensions and len(words), a relationship you don't comment on.

        grid[row][col] = words[i]

Using column-major order would have led to a more natural grid[x][y], or grid[col][row] notation.

def print_words(words):

Please add a little to the signature:

def print_words(words, k=3):

The the magic number 3 in expressions like // 3 and % 3 becomes the parameter k.

# special case
if word_count <= 3:
    grid.append(['']*3)

This suggests that your rows assignment was unfortunate. This case is no longer "special" if you assign:

rows = (word_count + 2) // 3

(or (word_count + k - 1) // k)

The subsequent comparisons of rows with 2 also seem gratuitously "special", making it hard to believe they are correct for all inputs.

# special case
if 3 < word_count < 6:
    print (grid[0][0]+ ' '+grid[0][2]+' '+grid[1][1])
    print (grid[0][1]+ ' '+grid[1][0]+' '+grid[1][2])

Promoting words into new rows in that way suggests the original assignment was faulty. Without a written spec, without a docstring on the assignment function, it is hard to pin blame and hard to suggest a remedy. Write a comment, and declare the code buggy if it does not conform to the comment.

# print grid

Good, these nested loops are clear, transparent code.

        print (grid[row][col], end=' ')

Using .ljust() here would let you accommodate words of variable width.

Answer 3

Here's my attempt at finding a straightforward and concise solution.

• It uses a function to calculate the height of each column.
• It creates a rectangular table (a list of lists), with d columns and as many rows as needed (the height of the first column).
• It iterates over every column and row and fills the characters.

---

l = ['a','b','c','d','e','f','g','h','i','j']

def height(n, d, i=0):
    """Return the height of the i-th column, for n elements in d columns."""
    rows, longer_columns = divmod(n, d)
    return range(rows + (i < longer_columns))

def distribute_elements(words, d=3):
    '''Distribute words into a table of d columns as evenly as possible.'''
    n = len(words)
    table = [[''] * d for _ in height(n, d)]
    word = iter(words)
    for i in range(d):
        for j in height(n, d, i):
            table[j][i] = next(word)
    return table

for rows in distribute_elements(l):
    print(' '.join(rows))

It outputs:

a e h
b f i
c g j
d