蝦J8亂寫
最後zip unpack那邊
因為空list會出錯 搞了好久==
就先這樣了
應該很多地方可以優化
def survivedRobotsHealths(self, positions: List[int], healths: List[int],
directions: str) -> List[int]:
n = len(positions)
idx = [i+1 for i in range(n)]
positions, healths, directions, idx = zip(*sorted(zip(positions, healths,
directions, idx)))
positions, healths, directions, idx = list(positions), list(healths),
list(directions), list(idx)
print(positions)
stk = []
ans = []
for i in range(n):
if directions[i] == "R":
stk.append(i)
else:
while stk:
if healths[stk[-1]] < healths[i]:
stk.pop()
healths[i] -= 1
elif healths[stk[-1]] == healths[i]:
stk.pop()
healths[i] = 0
break
else:
healths[stk[-1]] -= 1
break
if not stk and healths[i] > 0:
ans.append(i)
ans = ans + stk
ans = [(idx[i], healths[i]) for i in ans]
if ans:
_, ans = zip(*sorted(ans))
return list(ans)
else:
return []