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.

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")