[Codility] Lesson5. PassingCars (Python)

1.문제

app.codility.com/demo/results/training7U86YM-SXN/

Test results - Codility

2.풀이

prefix sum을 응용해서 뒤에서부터 합계를 sum 리스트에 저장했다.

그 이유는 출발 index = 0이라면 0부터~끝까지 1의 개수 , index = 1 이라면 1부터~끝까지 1의 개수를 미리 저장하기 위해서이다.

1의 개수를 모두 sum하면 정답!

이때 주의할 것은 정답이 1000000000 를 초과하면 -1을 리턴해야한다. 

 

3.코드

# 0 (east) -> 1(west)
# returns the number of pairs of passing cars.
# num of passing cars > 1000000000 -> -1 return

def solution(A):
    sum = [0]*len(A)

    for i in range(len(A)-1, -1, -1):
        if i == len(A)-1:
            sum[i] = A[i]
        else:
            sum[i] = sum[i+1]+A[i]
    
    answer = 0
    for i in range(len(A)):
        if A[i] == 0:
            answer += sum[i]
        if answer > 1000000000:
            return -1
    return answer