That looks about right. You can also use some list comprehensions instead of the for loops to do the same thing in a more compact way.
Something like this: