CSES standard solution for range query problem giving TLE

Posted on

Problem

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.

Solution

• 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]
``````