CodeChef - Red Alert

Question

I have solved this problem from the CodeChef platform using the Python prorgamming language

Here is the code:

test_cases = int(input())
for _ in range(test_cases):
 n, d, h = [int(x) for x in input().split()]
 rain_amounts = [int(y) for y in input().split()]
 red_alert = False 
 water_level = 0
 for rain_amount in rain_amounts:
 if rain_amount > 0:
 water_level += rain_amount 
 if rain_amount == 0:
 if water_level < d:
 water_level = 0
 else:
 water_level -= d
 if water_level > h:
 red_alert = True 
 if red_alert:
 print("YES")
 else:
 print("NO")

My solution received an ACCEPTED judgement and I would like feedback on my implementation, coding style, and everything in between.

Thank you.

• N3buchadnezzar: Please make the question self contained. E.g include the information from the link =)

Answer 1

I like your code =) What it does is clear, and your variable names and spacing is correct. There is not much of an algorithm to speak of here, so instead I will focus on structure.

It is always a good idea to separate inputs from actions. This can be done using a if_main guard, and clear functions. Something like this passes the first test case

def is_red_alert(precipitation, drainage, waterlogging):
 water_level = 0
 for rain_today in precipitation:
 if rain_today > 0:
 water_level += rain_today
 if rain_today == 0:
 if rain_today < drainage:
 water_level = 0
 else:
 water_level -= drainage
 return water_level > waterlogging
if __name__ == "__main__":
 test_cases = int(input())
 for _ in range(test_cases):
 _, drainage, waterlogging = [int(x) for x in input().split()]
 precipitation = [int(y) for y in input().split()]
 if is_red_alert(precipitation, drainage, waterlogging):
 print("YES")
 else:
 print("NO")

Otherwise the code is the same. I will change the name of a few of the variables throughout as naming is hard.. =) We can tidy up the logic a bit using min and max. Something like this

for day in precipitation:
 accumulation += max(0, day)
 if day == 0:
 accumulation -= min(drainage, accumulation)
 elif accumulation > limit:
 return True
return accumulation > limit

Firstly we do not have to check if the precipitation of the current day is negative. We can just use max(0, today) [Yes, this is slower]. Instead of setting the water level (or in my case the accumulation of water) we can just subtract the smallest number between drainage and accumulation. I will let you go through to check out what happens when drainage < accumulation, and drainage > accumulation. Similarly we only need to check if we are above the limit if the rain today is greater than zero. Otherwise there is no change.

We can also add some descriptive typing hints to further explain what is going on in the code

Code

from typing import Annotated
Precipitation = Annotated[
 list[int],
 "The amount water released from clouds in the form of rain, snow, or hail over n days",
]
Drainage = Annotated[
 int,
 "The amount of water removed on a dry day",
]
WaterMax = Annotated[
 int,
 "The total amount of water before triggering an warning",
]
def is_red_alert(
 precipitation: Precipitation, drainage: Drainage, limit: WaterMax
) -> bool:
 "Returns True if the citys water level becomes larger than a given limit"
 accumulation = 0
 for day in precipitation:
 accumulation += max(0, day)
 if day == 0:
 accumulation -= min(drainage, accumulation)
 elif accumulation > limit:
 return True
 return accumulation > limit
if __name__ == "__main__":
 test_cases = int(input())
 for _ in range(test_cases):
 _, drainage, water_limit = list(map(int, input().split()))
 precipitation = list(map(int, input().split()))
 if is_red_alert(precipitation, drainage, water_limit):
 print("YES")
 else:
 print("NO")

• King Cold: Thanks for such a comprehensive and thoughtful review.