[BOJ] 우체국2

[BOJ] 우체국2

Problem URL : 우체국2

Comments

처음에 다음과 같이 풀었다가 시간 초과(TLE)가 났다.

V, P, L = map(int, input().split())
towns = list(map(int, input().split()))


dp = [[[1e9] * V for _ in range(V)] for _ in range(P + 1)]
dp_offices = [[[[]] * V for _ in range(V)] for _ in range(P + 1)]
dist = [[0] * V for _ in range(V)]

for i in range(V):
    for j in range(V):
        if i + j >= V:
            dist[i][i+j-V] = L + towns[i+j-V] - towns[i]
        else:
            dist[i][i+j] = towns[i+j] - towns[i]

for i in range(V):
    for j in range(V):
        length = (j - i + V) % V + 1
        office = i
        dist_sum = 0
        for k in range(length):
            town = i + k
            if town <= office:
                dist_sum += dist[town % V][office % V]
            else:
                dist_sum += L - dist[town % V][office % V]
        dp[1][i][j] = dist_sum
        dp_offices[1][i][j] = [office % V]

        for k in range(1, length):
            office = i + k
            dist_sum += (k * 2 - length) * dist[(office-1) % V][office % V]
            if dp[1][i][j] > dist_sum:
                dp[1][i][j] = dist_sum
                dp_offices[1][i][j] = [office % V]


for p in range(2, P+1):
    for i in range(V):
        j = i + p - 1
        j %= V
        while j != i:
            mid = i
            while (mid+p-1) % V != j:
                tmp = dp[1][i][mid % V] + dp[p-1][(mid + 1) % V][j]
                if dp[p][i][j] > tmp:
                    dp[p][i][j] = tmp
                    dp_offices[p][i][j] = dp_offices[1][i][mid % V] + dp_offices[p-1][(mid+1) % V][j]
                mid += 1
            j += 1
            j %= V

ans = 1e9
for i in range(V):
    j = (i + V - 1) % V
    if ans > dp[P][i][j]:
        ans = dp[P][i][j]
        ans_offices = dp_offices[P][i][j]
ans_offices.sort()
print(ans)

for i in range(len(ans_offices)):
    print(towns[ans_offices[i]], end=' ')