Sorry, I came 3 years later. Buy I want to poin out that the problem ask for a function that take an array.

Touching anything in main is contrary to the excercise spirit. The program have an array in some place. The excercise main is there to make a mashed input and to test the answer. In python are better solutions to test a function ( or a class or a completet lib by the way) but this is the aproach used here.

So, you can do something in the line you do. For example you can make it whit iterators that take care of reading the file until needed and storing only the numbers you need and consuming them, that would be an interesting answer, but not what the problem is about.