백준4158🐨CD - 집념의 이터레이션 feat. map은 메모리를 먹는다.

https://www.acmicpc.net/problem/4158

 

테스트 케이스가 여러개인데 하나라고 착각하고 이상하게 짰습니다.
(문제 잘 안 읽는다고 아이들에게 주의 주는데... 반성합니다.)

n,m,*nums = map(int,open(0).read().split())
print(len(set(nums[:m])&set(nums[m:-2])))

그래서 예전에 성공한 코드를 보았습니다.

from sys import stdin

i = 0
data = stdin.readlines()
while True:
    m, n = map(int, data[i].split())
    if m == 0 and n == 0:
        break
    else:
        i+=1
        print(len(set(data[i:i+m]) & set(data[i+m:i+m+n])))
        i+=n+m

i를 따로 관리하면서 인덱스로부터  i~m, i+m~ i+m+n까지 슬라이싱 해서 집합 만들어서 세는 코드입니다.

i를 따로 관리하는 것 부터 마음에 안 드는 것 투성이라... 다음처럼 짜 보았습니다.

from itertools import islice

arr = iter(map(int,open(0).read().split()))
while True:
    m, n = next(arr), next(arr)
    if m & n:
        print(len(set(islice(arr, m)) & set(islice(arr, n))))    
    else:break

이터레이터를 써서 풀려고 해 보았습니다. 그런데 메모리 초과가 나왔습니다.

from itertools import islice

arr = iter(open(0).read().split())
while True:
    m, n = int(next(arr)), int(next(arr))
    if m & n:
        print(len(set(islice(arr, m)) & set(islice(arr, n))))    
    else:break

혹시나 하고 map을 풀어 보았더니 메모리 초과가 되지 않았습니다.
1708ms입니다.

map을 커대한 리스트에 돌리면 메모릴 아주 많이 먹습니다....
라는 것을 배웠습니다!!!

그런데 가만히 생각해 보니 매번 주어지는 m과 n을 잘 활용하지 못한다는 느낌이 들었습니다.

그리고 번뜩이는 아이디어!

from itertools import islice

arr = iter(open(0).read().split())
with open(1, 'w') as f:
    while True:
        mn = int(next(arr)) + int(next(arr))
        if mn:
            f.write(str(mn-len(set(islice(arr, mn))))+'\n')
        else:break

처음부터 m과 n을 불러와서 mn으로 합쳐(더해)버립니다.
이로서 변수 할당이 1 줄었습니다.
조건 검사가 m&n에서 mn으로 1개 줄었습니다.

집합을 각각 만들지 않고 M+N길이만큼 잘라서 한 번에 만듭니다. 이 때, 중복되는 원소만큼 줄어듭니다.
슬라이싱 1번, 셋 함수 1번 &연산자 한 번이 줄었습니다.

이미 구해 놓은 mn(m+n)에서 위 집합의 길이를 빼면 교집합의 개수가 나옵니다!
야호!

 

---

이 문제를 리스트로도 다시 풀어 보았는데. 이 코드는 ...
파이썬 3에서 1592ms 파이파이에서는 메모리 초과가 나왔다!!! 같은 코드인데!!

arr = list(open(0).read().split())
i = 2
res = []
while True:
    mn = int(arr[i - 2]) + int(arr[i - 1])
    if mn:
        res.append(str(mn - len(set(arr[i:i + mn]))) + '\n')
        i += mn + 2
    else:
        break
open(1, 'w').write("".join(res))

 

최종적으로 

해당 코드는 아래와 같다.

from itertools import islice

arr = iter(open(0).read().split())
while True:
    mn = int(next(arr)) + int(next(arr))
    if mn:
        print(mn-len(set(islice(arr, mn))))
    else:break

여기서 또 알 수 있었던 것은

반복 출력의 경우with open을 쓰지 않으면 에러가 나며(매번 close하면 될 것 같긴 하다.)

(1회당 출력이 많지 않으면)  open(1,'w')를 쓰는 것 보다 print가 빠를 수 있다.(with 의 영향 같다.)

하지만 파일을 한 번 열어두고 with로 하는 게 sys.stdout보다는 빨랐다.(import 오베헤드인 듯)