Leetcode 3629. Minimum Jumps to Reach End via Prime Teleportation

문제

You are given an integer array nums of length n.

You start at index 0, and your goal is to reach index n - 1.

From any index i, you may perform one of the following operations:

- Adjacent Step: Jump to index i + 1 or i - 1, if the index is within bounds.
- Prime Teleportation: If nums[i] is a prime number p, you may instantly jump to any index j != i such that nums[j] % p == 0.
Return the minimum number of jumps required to reach index n - 1.

길이가 n인 정수 배열 nums가 주어집니다.
당신은 인덱스 0에서 시작하며, 목표는 인덱스 n - 1에 도달하는 것입니다.
임의의 인덱스 i에서 다음 연산 중 하나를 수행할 수 있습니다:
인접한 위치로 점프: 인덱스가 범위 내에 있다면 i + 1 또는 i - 1로 점프합니다.
소수 순간이동: nums[i]가 소수 p라면, nums[j] % p == 0을 만족하는 다른 인덱스 j(j != i)로 즉시 점프할 수 있습니다.
인덱스 n - 1에 도달하는 데 필요한 최소 점프 횟수를 반환하세요.

예시

입력: nums = [1, 2, 4, 6]
출력: 2
설명:
최적의 점프 순서 중 하나는 다음과 같습니다:
인덱스 i = 0에서 시작합니다. 인접한 위치인 인덱스 1로 점프합니다.
인덱스 i = 1에서 nums[1] = 2이며, 이는 소수입니다. 따라서 nums[3] = 6이 2로 나누어 떨어지므로 인덱스 i = 3으로 순간이동합니다.
그러므로 정답은 2입니다.

제한

• $1 <= n == nums.length <= 10^5$
• $1 <= nums[i] <= 10^6$

풀이

해당 문제는 특정 정점에서 특정 정점으로 이동할수 있는것이 쉽게 예측 혹은 측정되기 어렵고, 갯수도 파악하기 어렵기 때문에 DP가 아닌 BFS를 사용하는것이 좋다.

나는 문제를 쪼개어서 풀었는데, 여기 쓰인 트릭은
Leetcode 실행 시간 줄이는 방법
해당 글 1번에서 확인할수 있다.

기본 생각

• 문제를 보면 알다시피 양옆도 이동할수 있지만 소수에 한정해, 자신을 소인수로 가지는 다른 숫자로도 이동할수 있다.
• 즉 {소수} -> {해당 소수를 소인수로 가지는 수} 라는것이다.
• 나는 여기서 약간 역발상을 통해, 이를 거꾸로 뒤집으면
• {소수를 소인수로 가지는 수} -> {소수} 라고 생각할수 있다고 봤다.
• 특정 소수를 소인수로 가지는 수를 찾는것 보다는 소수를 찾는게 더 쉬우니까 이방식을 써보기로 했다.

그래서

-- Class 밖 --
A. 에라스토테네스의 체를 이용해 Smallest Prime Factor와 소수 여부를 저장해둔다.
B. 구한 SPF를 이용해 소인수분해를 하고, 소수의 목록을 리턴하는 함수를 만든다.
-- Class 바깥 연산 끝 --

BFS의 준비

1. 배열에 있는 모든 소수의 위치를 구해서 저장한다.
2. 맨 뒤에 있는 값으로 부터 시작하는 BFS를 하나 만든다.
3. 이 외에 다익스트라 처럼 최소거리를 측정해서 저장하는 배열과 사용한 소수를 저장하는 set을 준비한다.

• 사용한 소수를 정리하는 이유는, 도착 지점이 특정 소수이기 때문에 한번 해당 소수로 타고 들어갔다면 다른 숫자가 굳이 거기 다시 갈 필요가 없기 때문.

BFS

1. 갈수 있는 소수들과 양옆 숫자를 q에 추가한다.
2. 다익스트라처럼 계속 이어 나간다.
3. 0에 도착하면 거리를 리턴한다.

MAX_NUM = 10 ** 6

spf = [i for i in range(MAX_NUM + 1)]
prime = bytearray(b'\x01') * (MAX_NUM + 1)
prime[0] = 0
prime[1] = 0

for i in range(2, int(MAX_NUM ** 0.5) + 1):
    if not prime[i]:
        continue

    for j in range(i * i, MAX_NUM + 1, i):
        if spf[j] == j:
            spf[j] = i

        prime[j] = 0

@lru_cache(None)
def prime_factors(target: int) -> List[int]:
    if target < 2:
        return []

    ans = []
    while target >= 2:
        f = spf[target]
        ans.append(f)
        while target % f == 0:
            target //= f

    return ans

class Solution:
    def minJumps(self, nums: List[int]) -> int:

        N = len(nums)

        prime_position = defaultdict(list)
        for i, v in enumerate(nums):
            if prime[v]:
                prime_position[v].append(i)

        check = [float('inf')] * N
        check[-1] = 0
        q = deque([(N - 1, 0)])

        used_factor = set()

        while q:
            pos, dist = q.popleft()

            if pos == 0:
                return dist

            goto = set()

            for factor in prime_factors(nums[pos]):
                if factor in used_factor:
                    continue
                used_factor.add(factor)

                if factor not in prime_position:
                    continue

                for position in prime_position[factor]:
                    goto.add(position)

            if pos < N - 1:
                goto.add(pos + 1)
            if pos > 0:
                goto.add(pos - 1)

            for dest in goto:
                if check[dest] <= dist + 1:
                    continue

                check[dest] = dist + 1
                q.append((dest, dist + 1))

        return -1