# 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 =  # 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 = 
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 = 
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]
``````