2D array inquiries

I am trying to solve a problem as described below:

You are given an input of a n*m forest.

You have to answer a group of inquiries. Your task is to return the number of trees for every inquiry.

Input: In the first line you are given 3 integers; n,m and q. This means that the size of the forest is n×m squares and the number of inquiries is q.

Then you get n rows that describe the forest. Every square is either blank (.) or a tree (*).

Then you are given q lines of inquiries. Every line has 4 integers y1, x1, y2 and x2, that define the area.

Output: Return the number of trees in the inquired area.

Limits:

1≤n,m≤1000

1≤q≤10^5

1≤y1≤y2≤n

1≤x1≤x2≤m

Example:

Input:

   4 7 3

.*...*.

**..*.*

.**.*..

......*

1 1 2 2

3 2 4 7

2 5 2 5


Output:

   3

4

1


My code is below; first it initializes the array so that every square is initialized to the number of trees in the rectangle from the top left square to the square. That takes O(n^2). After that I am able to answer the inquiries in O(1).

firstline = input().split()

rownumber = int(firstline[0]) #How many rows
rowlength = int(firstline[1]) #How many columns
numberOfInquiries = int(firstline[2])

forest = []

def printOut(someList):
for x in range(0, len(someList)):
print(someList[x])
#This function exist just for debugging

for x in range(0, rownumber):
forest.append([])
UI = list(input())
if UI[0] == ".":
first = 0
else:
first = 1
forest[x].append(first)
for y in range(1, rowlength):
if UI[y] == "*":
forest[x].append(int(forest[x][y-1])+1)
else:
forest[x].append(int(forest[x][y-1]))

for column in range(0, rowlength):
for row in range(1, rownumber):
forest[row][column] += forest[row-1][column]

print("")
#Everything is initialized

for x in range(0, numberOfInquiries):
UI = input().split()
y1 = int(UI[0])-1
x1 = int(UI[1])-1
y2 = int(UI[2])-1
x2 = int(UI[3])-1
first = forest[y2][x2]
second = forest[y1-1][x2]
third = forest[y2][x1-1]
fourth = forest[y1-1][x1-1]
if x1 == 0:
third = 0
fourth = 0
if y1 == 0:
second = 0
fourth = 0
print(first-second-third+fourth)


My question remains: how can this be optimized? The initialization takes 0.7-0.8 s in the larger cases and answering the inquiries also takes about the same. This is about 50% too slow for the largest test cases (time limit =1s). Are those just too much for Python? Most of the people solving these problems use C++.

• What challenge is this? – Peilonrayz May 8 '18 at 8:06
• A training problem from the Finnish Olympiad in Informatics page. cses.fi/problemset. This problem is only available in Finnish, though. – Kurns May 8 '18 at 14:31

2 Answers

1. You can use map or a list comprehension to allow placing assignment to rownumber, rowlength and numberOfInquiries on the same line:

rownumber, rowlength, numberOfInquiries = map(int, input().split())

2. When using range you don't have to specify 0. range(rownumber) works fine.

3. I find it hard to understand both reading from input and mutating it at the same time. You may want to build the forest of input before doing any mutations.

forrest = (
input()
for _ in range(rownumber)
)

4. You can abuse Booleans as they subclass ints, issubclass(bool, int), and so you can just use first = UI[0] == '*'.

5. You can use itertools.accumulate, so that you don't have to perform your y loop.
6. If you swap your loops when adding the previous row to the next you can use the pairwise recipe.

for prev, curr in pairwise(forest):
for column in range(rowlength):
curr[column] += prev[column]

7. Follow PEP 8, and use snake_case, rather than TitleCase or lowercase.

And so I'd use:

import itertools

row_number, row_length, number_of_inquiries = map(int, input().split())

def pairwise(iterable):
"s -> (s0,s1), (s1,s2), (s2, s3), ..."
a, b = itertools.tee(iterable)
next(b, None)
return zip(a, b)

forest = [
list(itertools.accumulate(
item == "*" for item in row
))
for row in (
input()
for _ in range(row_number)
)
]

for prev, curr in pairwise(forest):
for column in range(row_length):
curr[column] += prev[column]

for x in range(number_of_inquiries):
y1, x1, y2, x2 = (int(i) - 1 for i in input().split())
first = forest[y2][x2]
second = forest[y1-1][x2]
third = forest[y2][x1-1]
fourth = forest[y1-1][x1-1]
if x1 == 0:
third = 0
fourth = 0
if y1 == 0:
second = 0
fourth = 0
print(first-second-third+fourth)


Since you need to squeeze another 0.2 seconds, you could try setting the flush argument to print to False, to see if flushing is the limiting factor. This would look like:

print(first-second-third+fourth, flush=False)


Alternately you can do this yourself by not calling print.

ret = ''
for x in range(number_of_inquiries):
y1, x1, y2, x2 = (int(i) - 1 for i in input(ret).split())
# ...
ret = first-second-third+fourth

• Nope, still exceeding the time limit. Thanks for all your helpful tips, though. – Kurns May 8 '18 at 16:14
• You cut down the time by about 0.5 seconds. Just needs 0.2 off, it is very close! – Kurns May 8 '18 at 16:19
• @Kurns you could probably do that by using input as print. ret = ''; for x in ... input(ret) ... ret = first-second-third+fourth – Peilonrayz May 8 '18 at 16:21
• Could you elaborate? I don't really understand. – Kurns May 8 '18 at 16:41
• @Kurns I've added what I mean to my answer, if this isn't fast enough, then I suggest that you add a profile of your code. This means that I won't be guessing at what the problem is, and will be able to focus on making what is slow faster – Peilonrayz May 8 '18 at 17:30

I like the pre-initialisation approach.

One thing that you should look into is replacing explicit lists and append with approaches that operate on lists directly. List or generator comprehensions should be your go to tool for many things.

For example looking at this block of code

    UI = list(input())
if UI[0] == ".":
first = 0
else:
first = 1
forest[x].append(first)
for y in range(1, rowlength):
if UI[y] == "*":
forest[x].append(int(forest[x][y-1])+1)
else:
forest[x].append(int(forest[x][y-1]))


We can map the dots and stars to zero and one:

UI = list(input())
forest_line = [1 if x == "*" else 0 for x in UI]


Python has a useful bunch of built in functions in the itertools module, and in particular itertools.accumulate gives the prefix sums in a list. You'd need to import it at the start of the script.

forest[x] = accumulate(forest_line)


We can even collapse all this together. Here nothing gets assigned to a list until the very end, which makes things faster by avoiding storing the intermediate values.

forest[x] = list(accumulate(1 if x == "*" else 0 for x in input()))


You will generally find that this style of coding works faster than explicitely looping and appending, and is considered more readable by python enthusiasts. (Although there is a limit to how much you can cram into one line before it becomes too much of a labyrinth)