Read millions of integers from stdin and process them in Python within 2 seconds

Question

This is a competitive programming problem which has never been solved in Python within the time limit of 2 seconds. It is about Disjoint Set Union:

Source: Algorithms In Context #10: Disjoint Sets

Description:

The input comes from a text file with 100 test cases defined by a few million of numbers. The first line of each test case indicates:

• number of nodes (2 < N < 10_000)
• number of links (0 < L < 100_000)
• number of queries (1 < Q < 1_000)

The subsequent numbers are shown in pairs (two per line), ranging from 1 to N, indicating L links and then the Q queries. This is a sample test case:

5 3 2
1 2
2 3
4 5
1 3
3 4

The goal is to respond each query with Y or N indicating if the queried pair belongs to the same group or not. For the above example, the expected output is:

Y
N

Because the links 1-2-3 bring 1 and 3 to the same group, but 3 and 4 are still in distinct groups.

The first 90 test cases are not so big (1000 200 100) while the last 10 test cases use the maximum size (10000 100000 1000).

Solution:

I have submitted hundredths of different versions, but my fastest try solved just 97 of the 100 test cases. The conventional solution involves building tree-like structures. And I did test all kinds of that. But the flood-fill approach ended up doing better (resolved more test cases). There you go:

from sys import stdout
from sys import stdin
from array import array

write = stdout.write

arrBase = array('h', (0 for _ in range(10001)))

# Stores a number indicating group ID of each
# node (0 indicates it hasn't been set)
arr = []

# Stores a Set() containing the direct
# link relations (graph edges) for each node
rels = []

def floodfill(groupid):
    ni = 0
    ids = set()
    ids.add(groupid)
    nodes = [groupid]
    while ni < len(nodes):
        n = nodes[ni]
        ni += 1
        if arr[n] != groupid:
            arr[n] = groupid
            newIds = list(rels[n].difference(ids))
            ids.update(newIds)
            nodes.extend(newIds)

nodes = 0
links = 0
queries = 0
q = 0 # query counter

outbuffBase = array('u', ('N' for _ in range(1000)))
outbuff = []

# flag to assert a blank line is printed between each test case
startWithLineBreak = 0

for line in map(bytes.split, stdin.buffer):
    if q == queries:
        nodes, links, queries = map(int, line)

        # initializing test case
        arr = arrBase[:nodes+1]
        rels = [set() for _ in range(nodes + 1)]
        outbuff = outbuffBase[:queries]
        q = 0
        continue

    a, b = map(int, line)

    # storing a link relation
    if links > 0:
        links -= 1
        rels[a].add(b)
        rels[b].add(a)

    # executing a query
    else:
        va = arr[a]
        if va == 0:
            floodfill(a)
            va = a
        vb = arr[b]
        if va == vb:
            outbuff[q] = 'Y'
        q += 1
        if q == queries:
            write("\n" * startWithLineBreak + "\n".join(outbuff[:queries]) + "\n")
            startWithLineBreak = 1

stdout.flush()

Question: is there any performance optimization which can be done in this code?

Notes:

• We can't use external libraries. Numpy arrays actually did make the code slower anyway.
• Performance measurements have shown that the program takes more than 1 second just to read the numbers, so the bottleneck might be in the split method. The second slowest process is the actual flood-fill. While the initialization of variables and the output printing are quite fast (no improvement needed).

Answer 1

use isort

You may as well, there's no reason not to. This is phrased a bit oddly, on two lines rather than one.

from sys import stdout
from sys import stdin

Run isort and then don't give it a second thought. There's bigger fish to fry.

write = stdout.write

That's weird. You don't like print() ? If there is e.g. some performance benefit or other advantage, then call it out, in a # comment.

pythonIsNotJava

arrBase = array('h', (0 for _ in range(10001)))

Pep8 asked you nicely to call it arr_base, 'nuff said.

That for loop involves interpreting a bunch of bytecode n times. I wonder if some variation on (0,) * n or [0] * n would win, here, so compiled C code is doing the nitty gritty.

class

This whole codebase would be a natural fit for some new class you define.

meaningful identifier

arr = []
rels = []

Both of those could potentially be good local variable names. As it is, you have them at module scope (global), so the documentation burden is greater. And, I have no idea at this point what they mean; a # comment would have been helpful. Or pick more meaningful names, so no head scratching and no comment is necessary.

Maybe it should be named neighbors ? I'm not sure how to mentally pronounce rels.

EDIT:

OIC, pronounce it "edge", where Author intended "link relation".

set literal

    ids = set()
    ids.add(groupid)
    nodes = [groupid]

As those first two lines, prefer ids = {groupid}, for parallel construction with the third line.

tuple unpack

        nodes, links, queries = map(int, line)

This looks lovely. Thank you for the idiomatic unpack, it's very clear.

I am a little worried that line is not yet decoded from bytes to the more customary str, but I suppose the int mapping dealt with it for us.

defaultdict

        rels = [set() for _ in range(nodes + 1)]

Maybe assigning defaultdict(set) would be helpful here?

automated tests

This submission contains no test suite. It would be stronger if it provided one.

A test suite is a tool for educating newly on-boarded engineers about how to call the target code and what results to expect. But it is also a starting point for benchmarking efforts.

appropriate signature

def floodfill(groupid):

This is a bit vague, as we're messing with various globals.

It would be better to explictly call them out:

def floodfill(groupid, arr, rels):

That gives the Gentle Reader a heads up that we may be reading or mutating some large datastructures.

It would be much better if the function included a """docstring""", explaining what it accepts as input and what cool thing it has accomplished upon returning. For example, if I were to write a unit test that called floodfill(), how would I know that the Right Thing had happened? What should I assert? What is the promised post-condition?

bool

outbuffBase = array('u', ('N' for _ in range(1000)))

Prefer to store {0, 1} bytes instead of utf32 characters, for 75% memory savings.

functions

Use them.

If you had expressed your logic with named def functions, which included """docstrings""", it would be much easier to understand the contracts, understand where the CPU cycles went, and propose refactorings of "same contract, faster implementation".

---

This codebase achieves a subset of its design goals.

I would not be willing to delegate or accept maintenance tasks on it in its current form.