Parsing input
input().strip().split()
can be reduced to input().split()
as split
without parameters has extra logic to remove empty strings in the result. It feels also more idiomatic to unpack the two integers using r, c = map(int, input().split())
Parsing princess position
First of, you could create the grid
using a list-comprehension rather than append
:
grid = [input() for _ in range(n)]
Note that I dropped 0
from the range
parameter as it the default starting point.
But you don't even need to build the whole grid since you’re only interested in knowing the princess position. You could search for it while parsing the input:
def princess_position(grid_size):
for row in range(grid_size):
try:
return row, input().index('p')
except ValueError:
# Princess not on this line
pass
# Princess not in grid
return -1, -1
grid_size = int(input())
bot_row, bot_column = map(int, input().split())
princess_row, princess_column = princess_position(grid_size)
print(next_move(bot_row, bot_column, princess_row, princess_column))
Comparing positions
You don't need to put parenthesis around ternary operators. You can also simplify a - b > 0
to a > b
; same for other comparisons. Last thing, you can perform less work by comparing a == r
first and returning left or right before comparing for up and down:
def next_move(bot_row, bot_column, princess_row, princess_column):
if bot_row == princess_row:
return 'LEFT' if bot_column > princess_column else 'RIGHT'
return 'UP' if bot_row < princess_row else 'DOWN'
Top-level code
It is better practice to wrap your top-level code under an if __name__ == '__main__'
statement, as it allows you to import your file without executing the whole thing:
def next_move(bot_row, bot_column, princess_row, princess_column):
if bot_row == princess_row:
return 'LEFT' if bot_column > princess_column else 'RIGHT'
return 'UP' if bot_row < princess_row else 'DOWN'
def princess_position(grid_size):
for row in range(grid_size):
try:
column = input().index('p')
except ValueError:
# Princess not on this line
pass
else:
return row, column
# Princess not in grid
return -1, -1
if __name__ == '__main__':
grid_size = int(input())
print(next_move(
*map(int, input().split())
*princess_position(grid_size))
I also changed the way to use unpacking to simplify variable managements.