[Codility] Passing Cars (Python3)

https://app.codility.com/programmers/lessons/5-prefix_sums/passing_cars/

교차하는 차의 수를 구하는 문제.
리스트 A의 원소(자동차)들은 인덱스 순서로 양 방향으로 지나간다.
자동차 하나가 한 쪽으로 지나간 이후에 나오는 다른 차들이 반대방향으로 가는 경우, 즉 교차되는 경우의 수가 몇 개인지 찾는 문제이다.

이 문제의 핵심은 0뒤에 나오는 1의 총 개수를 찾는 문제이다.

A가 가질 수 있는 원소의 수가 10만개이기 때문에 O(N^2) 미만의 시간 복잡도로 문제를 풀어야 했다. 이중 포문을 돌려서 하나씩 찾으면 안되기에 어떤 방법이 있나 한시간을 넘게 고민했다.

그러다가 얼마 전에 풀었던 Max Count 문제가 떠올랐다.
이 문제도 이중 포문으로 쉽게 풀 수 있는 문제였지만 효율성을 위해서 특정 조건에서 count를 세주는 변수를 활용하여 문제를 풀었었다.

이번에도 비슷한 방법으로 zero_num이라는 변수를 활용했고 문제를 풀 수 있었다.

나의 풀이

핵심은 0뒤에 나오는 1의 개수를 모두 카운팅하는 것이다.
반복문을 돌면서 0의 개수를 세주고 1이 나왔을 때 현재까지의 0의 개수만큼 cnt 값을 증가시켜 문제를 해결하였다.

def solution(A):
    # write your code in Python 3.6

    cnt = 0
    zero_num = 0
  
    for i in A:
        if i == 0:
            zero_num += 1
        elif i == 1:
            cnt += zero_num

    if cnt > 1000000000:
        return -1
    else:
        return cnt

    pass