[Lua] 백준 19566 수열의 구간 평균

반딧불이 만든 문제고, 내가 다른 분의 풀이를 참고하지 않고 스스로 푼 문제들 중 난이도가 가장 높은 문제이다.

 

푼 지 꽤 오래되었고 푼 글도 올린 줄 알았는데 안 올라와있어서 지금 올린다.

 

이 문제를 풀려면 먼저 백준 2015번 수들의 합 4 문제의 풀이를 이해하고 있어야 한다.

 

n,k=io.read("*n","*n")
set={}
set[0]=1
sum,ans=0,0
for i=1,n do
    a=io.read("*n")
    sum=sum+a
    if set[sum-k]~=nil then ans=ans+set[sum-k] end
    if set[sum]==nil then set[sum]=1 else set[sum]=set[sum]+1 end
end
io.write(ans)

 

이 코드에서 sum 변수는 i번째 숫자까지의 합이다. 

 

sum[i]-sum[j]=k를 만족하는 경우의 수를 i,j에 대해 2중 반복문으로 돌리면 N^2 이기 때문에 시간 초과가 난다.

 

그래서 이 식을 변형해 sum[i]-k=sum[j] 로 만들고 i에 대해 반복문을 돌려서 O(N)으로 푸는 게 수들의 합 4 문제의 풀이이다.

 

set[0]=1인 이유는, sum[i]-k=0, 즉 sum[i]=k 인 경우(1번째부터 i번째까지 숫자의 합이 k인 경우)를 놓치지 않기 위함이다. 

 

sum[i]-k=sum[j]인 경우는 ans에 더해주기만 하면 된다.

 

그리고 set[sum[i]]를 업데이트 해줘야 한다. ( 예를 들어 sum[4]-k=sum[2] 일 때, 반복문이 i=2 에서 sum[2]인 경우의 수를 업데이트 안 해뒀으면 이 경우를 그냥 넘어가게 된다.) 그래야 임의의 수 x에 대해 set[x]가 누적 합이 x가 되는 경우의 수가 된다.

 

어쨌든 수들의 합 4에서의 핵심 식이 sum[i]-sum[j]=k 라면

 

이 문제에서의 핵심 식은 sum[i]-sum[j]=k*|i-j| 이다.

 

아까처럼 sum[i]-k*|i-j|=sum[j] 로 바꾸면 좌변에서 i,j 두 수를 알아야 하기 때문에 2중 반복문을 써야 한다.

당연히 이 문제도 O(N^2) 로 풀면 시간 초과가 난다.

 

이를 방지하려면 절댓값 기호를 풀어줘야 한다.

i>j 이면 sum[i]-ki=sum[j]-kj가 되고

i<j이면 sum[i]+ki=sum[j]+kj가 된다.

 

그리고 set[sum[i]-k]를 ans에 더하는 대신 set[sum[i]-ki]와 set[sum[i]+ki]를 모두 더해주도록 했다.

 

이 두 경우를 모두 세 주고, set[0]=1 처리도 똑같이 해 주는데, k가 0이면 i>j인 경우랑 i<j인 경우가 같아서 같은 경우를 2번 세는 문제가 발생한다. 그래서 i<j 인 경우를 ans에 더해줄 때는 k가 0이 아닐 때만 더하도록 코드를 짜서 AC를 받았다.

 

근데 나도 이 풀이가 정확한 건지는 잘 모르겠다...