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:

a disconnected graph with three components
Source: Algorithms In Context #10: Disjoint Sets


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:


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).


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()
    nodes = [groupid]
    while ni < len(nodes):
        n = nodes[ni]
        ni += 1
        if arr[n] != groupid:
            arr[n] = groupid
            newIds = list(rels[n].difference(ids))

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

    a, b = map(int, line)

    # storing a link relation
    if links > 0:
        links -= 1

    # executing a query
        va = arr[a]
        if va == 0:
            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


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


  • 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).
  • \$\begingroup\$ What kind of "tree-like structure" did you test? The way you phrased that worries me, because one of the best aspects of the Union-Find data structure is that it's an array (or two arrays depending on how you do it, but still: array). \$\endgroup\$
    – user555045
    Commented Jan 29 at 16:34
  • 1
    \$\begingroup\$ Yes @harold a tree structure represented in the form of an array of integers pointing to other indexes in the array, plus additional array with the tree size or rank (size of nodes below). Just like the links I added. I spent most of my time with this structures, but the flood-fill ended up being faster, because it's a O(n) solution were n is the number of nodes. It actually gets below O(n) since the flood fill is only calculated for the queried nodes. \$\endgroup\$ Commented Jan 29 at 19:59

1 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.


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.


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.


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

set literal

    ids = set()
    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.


        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?


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

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


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.

  • \$\begingroup\$ I have added the description for arr and rels variables. But my goal here is only to make the code run faster. This is why I used global parameters. Because that made the code faster. Notes: The input is being read in binary mode, skipping the decoding step to string. The first map statement is already splitting all numbers. The second one is only mapping sets of byte representing the numbers to actual python ints. I don't get why defaultdict was suggested. I used 'u' (utf32) character because that's a ready to output. I have tried also writing direct binary data to stdout.buffer.write. \$\endgroup\$ Commented Jan 29 at 7:15
  • \$\begingroup\$ Hey, joining stdin and stdout imports into a single import statement was enough to make the code go through. Thanks ! \$\endgroup\$ Commented Jan 29 at 7:39

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.