1621. Number of Sets of K Non-Overlapping Line Segments

[zeikar]

Problem Link

https://leetcode.com/problems/number-of-sets-of-k-non-overlapping-line-segments/

Problem Summary

1차원 위에 놓인 점 n개에서 서로 겹치지 않는 선분 k개를 그리는 경우의 수를 구하는 문제. 선분은 점을 둘 이상 덮어야 하고 끝점은 공유할 수 있다.

Solution

modular 나온 걸 보고 딱 DP겠거니 했다.

어제 2472처럼 겹치지 않게 k개를 고르는 문제인데 이번엔 최대 개수가 아니라 경우의 수다. 한 점에서 할 수 있는 건 이어붙이기, 끊고 넘어가기, 새로 시작하기 세 가지인데 이어붙이는 중인지를 상태로 넣어야 해서 3차원이 된다.

solve(i, j, c) = i번째 점을 볼 차례이고 선분을 j개 썼을 때의 경우의 수 (c는 이어붙이는 중인지)

c == 0: solve(i+1, j, 0) + solve(i+1, j+1, 1)
c == 1: solve(i+1, j, 1) + solve(i+1, j, 0) + solve(i+1, j+1, 1)

끝점을 공유할 수 있으니 c == 1에서 끊자마자 새로 시작하는 solve(i+1, j+1, 1)이 따로 필요하다. j == k이고 c == 0이면 남은 점은 그냥 두면 되니까 1을 돌려주면 된다.

처음엔 또 편하게 @cache로 짰다가 메모리 초과... 항목이 100만 개인데 인자 3개짜리 캐시는 항목당 130바이트 정도라 캐시만 125MB다. 3차원 배열로 바꾸니 통과.

시간복잡도는 O(nk)

참고로 수학적으로는 C(n + k - 1, 2k)로 바로 나온다. stars and bars로 증명된다. 자세한건 에디토리얼 참고.

Source Code

class Solution:
    def numberOfSets(self, n: int, k: int) -> int:

        dp = [[[ -1 for _ in range(2) ] for _ in range(k+1) ] for _ in range(n)]

        def solve(i, j, c):
            if j == k and c == 0:
                return 1
            if j > k:
                return 0
            if i == n:
                return 0

            if dp[i][j][c] != -1:
                return dp[i][j][c]
            
            if c == 1:
                dp[i][j][c] = (solve(i+1, j, 1) + solve(i+1, j+1, 1) + solve(i+1, j, 0)) % (10**9+7)
                return dp[i][j][c]
            else:
                dp[i][j][c] = (solve(i+1, j, 0) + solve(i+1, j+1, 1)) % (10**9+7)
                return dp[i][j][c]
        
        return solve(0, 0, 0)