A fast input and output template¶

~15 min · prerequisite: none inside this part; you should be comfortable with Parts 1 to 5

What you will be able to do¶

Read all of a judge's input at once and turn its words into integers.
Collect answers in a list and write them with one call at the end.
Organise a solution as a solve function that you can test with assert.

The idea¶

An online judge gives your program its input as text on standard input, which programmers call stdin. A test case can hold thousands or hundreds of thousands of numbers. The simplest reader is input(), which reads one line each time you call it. That is fine for a handful of lines. When a program calls input() once per number, each call adds overhead, and the overhead grows with the size of the test.

The template reads everything in one step instead. It takes the whole input, splits it into tokens, and then hands out the tokens one at a time. A token is a run of characters separated from the others by whitespace. Spaces and line breaks both count as whitespace, so the same code works whether the numbers sit on one line or on many.

Output follows the same idea. Printing inside a loop calls the output routine once per answer. The template instead builds the answer as a single piece of text, with line breaks between the values, and writes it once at the end.

The code has three parts. The function solve holds the problem logic. It takes an iterator, which is an object that gives out one item at a time, so it does not care where the tokens come from. The function main reads standard input and calls solve. The line if __name__ == "__main__": runs main only when you start the file as a program. When another program imports the file, main does not run, so you can test solve directly.

Worked example¶

This block reads standard input, so it cannot run in the browser runner. Save it as solution.py and run it on your own computer.

import sys


def solve(tokens):
    n = int(next(tokens))
    values = [int(next(tokens)) for _ in range(n)]
    return str(sum(values))


def main():
    data = sys.stdin.buffer.read().split()
    answer = solve(iter(data))
    sys.stdout.write(answer + "\n")


if __name__ == "__main__":
    main()

Line by line, this is what Python does:

import sys loads the sys module, which gives access to standard input and standard output.
def solve(tokens): defines a function that takes an iterator of words.
n = int(next(tokens)) asks the iterator for its next word, which is the count, and converts that word to an integer. The function next returns one item and moves the iterator forward.
The list comprehension asks for the next n words and converts each one to an integer.
return str(sum(values)) returns the answer as text, because the output must be text.
sys.stdin.buffer.read() reads all of standard input as bytes. .split() splits it at every run of whitespace, so you get a list such as [b"3", b"1", b"2", b"3"].
solve(iter(data)) turns the list into an iterator and runs the logic.
sys.stdout.write(answer + "\n") writes the answer and a line break in one call.
The last two lines call main only when the file is started directly.

If the file input.txt contains 3, then 1 2 3 on the next line, then python3 solution.py < input.txt prints 6.

Try it¶

The first block tests solve with text tokens and with byte tokens. The second shows that next returns items in order. The third tests a helper that joins answers with line breaks.

def solve(tokens):
    n = int(next(tokens))
    values = [int(next(tokens)) for _ in range(n)]
    return str(sum(values))


assert solve(iter("3 1 2 3".split())) == "6"
assert solve(iter("1 -5".split())) == "-5"
assert solve(iter(b"2 10 20".split())) == "30"

tokens = iter("2 5 6".split())
assert next(tokens) == "2"
assert next(tokens) == "5"
assert next(tokens) == "6"

def format_answers(answers):
    return "\n".join(str(a) for a in answers) + "\n"


assert format_answers([3, 7]) == "3\n7\n"
assert format_answers([10]) == "10\n"

Common mistakes¶

Comparing a token with a number before converting it. The first token is bytes, so this check fails:
```
tokens = b"5".split()
assert tokens[0] < 3
```
```
TypeError: '<' not supported between instances of 'bytes' and 'int'
```
Convert with int(tokens[0]) first.
Adding a number to a token that is still text. Writing n = next(tokens) and then n + 1 gives:
```
TypeError: can only concatenate str (not "int") to str
```
Wrap the token in int(...).
Mixing input() with a full read. After sys.stdin.buffer.read() has consumed everything, a later input() finds nothing and fails:
```
EOFError: EOF when reading a line
```
Passing numbers to "\n".join. The function join accepts only strings:
```
print("\n".join([1, 2]))
```
```
TypeError: sequence item 0: expected str instance, int found
```
Convert each answer with str(...) first.

Practice¶

Write solve(tokens) that reads one integer n and returns the text of n * n. Check that solve(iter(["4"])) returns "16".
Write solve(tokens) that reads an integer n, then n words, and returns how many of the words are "yes" as text. Check solve(iter("3 yes no yes".split())) returns "2".
Write solve(tokens) that reads an integer t, then t pairs of integers, and returns the sums on separate lines. Check solve(iter("2 1 2 3 4".split())) returns "3\n7".
Write solve(tokens) that reads n and then n integers and returns their maximum as text. Check solve(iter("4 3 9 2 5".split())) returns "9".
Stretch: write solve(tokens) that reads n and then n integers and returns the number of distinct values as text. Check solve(iter("4 1 1 2 2".split())) returns "2".

Solution

```python def solve(tokens): n = int(next(tokens)) return str(n * n)

assert solve(iter(["4"])) == "16" ```

```python def solve(tokens): n = int(next(tokens)) count = sum(1 for _ in range(n) if next(tokens) == "yes") return str(count)

assert solve(iter("3 yes no yes".split())) == "2" ```

```python def solve(tokens): t = int(next(tokens)) out = [] for _ in range(t): a = int(next(tokens)) b = int(next(tokens)) out.append(str(a + b)) return "\n".join(out)

assert solve(iter("2 1 2 3 4".split())) == "3\n7" ```

```python def solve(tokens): n = int(next(tokens)) values = [int(next(tokens)) for _ in range(n)] return str(max(values))

assert solve(iter("4 3 9 2 5".split())) == "9" ```

```python def solve(tokens): n = int(next(tokens)) values = [int(next(tokens)) for _ in range(n)] return str(len(set(values)))

assert solve(iter("4 1 1 2 2".split())) == "2" ```

Recap¶

Read all of standard input once with sys.stdin.buffer.read().split(), then take tokens one at a time.
Convert tokens with int before comparing or doing arithmetic on them.
Build the answer as text and write it once, and keep the solution in a solve function you can test with assert.