Graph Matrix with Numpy

There are so many code interview websites out there, but I read on LinkedIn about one called Code Signal and thought I would give it a try. In the "Arcade" mode, there is a section for graphs, which piqued my interest. If you aren't familiar with the power of Numpy for math with matrices, you'll be able to see below.

What is Numpy? It's a Python package that packs homogeneous data types into arrays in a form that is more condensed than heterogeneous Python lists; it breaks the work up into fragments for parallel processing; and it integrates C, C++, and Fortran into Python, which process faster than Python. (read about why Numpy is faster)

What is the problem? Code Signal sets up the scene with a king looking at improving traffic with one way roads between cities in his kingdom. Each city must have the same number of roads going in as going out. The connections between the cities is a list of lists, as seen in the following:

roadRegister = [[false, true,  false, false],
                [false, false, true,  false],
                [true,  false, false, true ],
                [false, false, true,  false]]

You can see in the above example that there is 1 True in the first list (city 1 => city 2). For the first city to pass the test, it would need only 1 True in the first column (position [0]. You can see that there is indeed 1 True connecting city 3 with city 1. That's 1:1. You can see that city 3 has 2 True values in position [0] and position [3], so it would need 2 True values in the same position column, which it does. This one would pass the test, so you would need to return True. How could you do that?

One way would be to iterate through the lists with two for-loops adding 1 for each True value using the index to keep track of where to add. As I said above, lists are less efficient than a Numpy array and Numpy can process much faster than a for-loop. So how do we execute this with Numpy?

import numpy as np
# import numpy and convert to an array specifying it's booleans
array = np.array(roadRegister, dtype=bool)

The array is the same list of booleans, but booleans are binary and can also be expressed as True = 1 and False = 0.
For all intents and purposes, these two matrices are equal.

[[0,1,0,0],             [[False, True, False, False],
 [0,0,1,0],          =   [False, False, True, False],
 [1,0,0,1],              [True, False, False, True],
 [0,0,1,0]]              [False, False, True, False]]

So with Numpy, you can take those booleans and add them up as if they were 0's and 1's.

# For the sum of the columns, axis = 0 (vertically)
col = np.sum(array, axis=0)

# For the sum of the rows/index, axis = 1 (horizontally)
row = np.sum(array, axis=1)

That's it. col and row will be arrays with 4 elements, the sums of the columns and rows respectively. That output for both is [1 1 2 1].

From there, you just need to see if the lists are the same and return True if they are and False if they aren't. Again, you could do a for-loop and iterate through each item and see if each item is equal; or you could solve it with another Numpy function array_equal(array1, array2) making your function magnitudes faster.

Here is the whole code that solves the puzzle fastest in Python.

import numpy as np
def solution(roadRegister):
    array = np.array(roadRegister, dtype=bool)
    col = np.sum(array, axis=0)
    row = np.sum(array, axis=1)
    return np.array_equal(row, col)