CSES standard solution for range query problem giving TLE

Problem statement is from CSES

This is a standard question where we are given a list of numbers and a number of queries. For each query, you have to give the sum of numbers in the given range.

My implementation in python:

def solve():
    n,q = [int(a) for a in input().split(' ')]
    arr = [int(a) for a in input().split(' ')]
 
    cumulativeArray = [0] # Initial entry added to make 1 based indexing
    sum = 0
    for i in range(len(arr)):
        sum += arr[i]
        cumulativeArray.append(sum)
 
    for i in range(q):
        x,y = [int(a) for a in input().split(' ')]
        print(cumulativeArray[y] - cumulativeArray[x-1])
 
solve()

This gives the required answers correctly for all my custom test cases. But gives TLE for huge test cases.

I believe this is happening due to split() method in my code.
Any leads are welcome.

To add on to @hjpotter92’s answer:

• If you’re not using i in for i in range(q), replace it with _, i.e. for _ in range(q) to indicate that you are not using the variable.
• It turns out the actual bottleneck is the many calls to print() in a for loop. If you tweak your code so it instead builds a string of all the answers to the queries in memory and prints it out in one go, it will pass all the test cases. I verified this by submitting the code below.

#!/usr/bin/env python3

n, q = map(int, input().split())
arr = map(int, input().split())

cumulative_sums = [0]
current_sum = 0
for num in arr:
    current_sum += num
    cumulative_sums.append(current_sum)

queries = (tuple(map(int, input().split())) for _ in range(q))
print(
    "n".join(
        f"{cumulative_sums[b] - cumulative_sums[a - 1]}" for a, b in queries
    )
)

You are already using the optimal solution algorithm for the problem. A few things, which might speed-up the computations (.split should not be the limiting factor in most likeliness)

Variable naming

Use names which match the problem description, so that it becomes easier to follow your code. So, instead of x, y it’d be a, b. Also, in python, it is a good practice to name your variables in lower_snake_case.

sum is an builtin function. Do not override this.

Computing len

You do not need to iterate over indices as the index is not really being used anyway. Lists can be iterated over their values itself.

Lazy iterators

Instead of generating list inline, you can use generators so that the values are yielded when iterating, and not consuming memory.

.split

Not passing a parameter to .split is same as providing whitespace as parameter.

Rewrite:

def solve():
    n, q = map(int, input().split())
    arr = map(int, input().split())
 
    cumulative_summation = [0]
    running_sum = 0
    # Not calculating `len` of `arr`, also, you don't need to calculate it.
    # the length is known to be `n`.
    for value in arr:
        running_sum += value
        cumulative_summation.append(running_sum)
 
    for i in range(q):
        a, b = map(int, input().split())
        print(cumulative_summation[b] - cumulative_summation[a - 1])

Your solution is alright and I got it accepted as-is. Just had to choose “PyPy3” instead of “CPython3”.

Anyway…

I’d use accumulate and a helper function to reduce code duplication and make the interesting code clearer:

from itertools import accumulate

def input_ints():
    return map(int, input().split())

def solve():
    _, q = input_ints()
    cumulative = [0, *accumulate(input_ints())]
    
    for _ in range(q):
        a, b = input_ints()
        print(cumulative[b] - cumulative[a - 1])

solve()

This still gets TLE with CPython3, though, and gets accepted in PyPy3 in the same 0.81 seconds as yours.

As @setris pointed out you can get it accepted by using a single print. Here’s my version of that, not mixing the calculations with the printing, which I find clearer. Got accepted with CPython3 in 0.69 seconds.

from itertools import accumulate

def input_ints():
    return map(int, input().split())

def solve():
    _, q = input_ints()
    cumulative = [0, *accumulate(input_ints())]

    results = []
    for _ in range(q):
        a, b = input_ints()
        results.append(cumulative[b] - cumulative[a - 1])

    print('n'.join(map(str, results)))

solve()

If you don’t mind the asymmetry of solve reading the input itself but not printing the output itself, you could also make it a generator:

from itertools import accumulate

def input_ints():
    return map(int, input().split())

def solve():
    _, q = input_ints()
    cumulative = [0, *accumulate(input_ints())]

    for _ in range(q):
        a, b = input_ints()
        yield cumulative[b] - cumulative[a - 1]

print('n'.join(map(str, solve())))

I do mind that asymmetry, though. Another symmetric way, instead of solve doing both input and output, would be to solve doing neither of them:

from itertools import accumulate

def input_ints():
    return map(int, input().split())

def solve(x, queries):
    cumulative = [0, *accumulate(x)]

    for a, b in queries:
        yield cumulative[b] - cumulative[a - 1]

_, q = input_ints()
x = input_ints()
queries = (input_ints() for _ in range(q))
results = solve(x, queries)
print('n'.join(map(str, results)))

Granted, now the “main block” doesn’t look nice. But solve is now very nice, and it can also be tested nicely by just calling it with data, for example:

>>> list(solve([3, 2, 4, 5, 1, 1, 5, 3],
               [[2, 4], [5, 6], [1, 8], [3, 3]]))
[11, 2, 24, 4]