Computer >> 컴퓨터 >  >> 프로그래밍 >> Python

파이썬으로 배열을 오른쪽으로 k번 회전한 후 i번째 요소 찾기

문제 이해하기

배열 nums와 값 k, 그리고 또 다른 값 i가 주어졌을 때, nums의 요소들을 오른쪽으로 k번 회전한 후 인덱스 i에 위치한 요소를 찾아야 합니다.

예를 들어 입력이 nums = [2,7,9,8,10], k = 3, i = 2라고 가정해 보겠습니다. 세 번 회전한 후 배열은 [9,8,10,2,7]이 되므로, 이때 i번째 요소는 nums[2] = 10이 됩니다. 따라서 출력값은 10입니다.

해결 접근 방법

이 문제는 다음 단계를 따라 해결할 수 있습니다.

  • 0부터 k까지 반복하면서 매 회전마다 다음을 수행합니다.
    • nums의 마지막 요소를 삭제(pop)하고, 그 요소를 인덱스 0 위치에 삽입(insert)합니다.
  • 모든 회전이 끝나면 nums[i]를 반환합니다.

구현 예제

아래 구현을 통해 더 잘 이해해 보겠습니다.

def solve(nums, k, i):
    for r in range(k):
        nums.insert(0, nums.pop())
    return nums[i]

nums = [2,7,9,8,10]
k = 3
i = 2
print(solve(nums, k, i))

입력

[2,7,9,8,10], 3, 2

출력

10

더 효율적인 개선 방법

위 방법의 시간 복잡도는 O(k × n)으로, k가 커지면 비효율적일 수 있습니다. 하지만 수학적 성질을 활용하면 배열을 실제로 회전하지 않고도 답을 바로 구할 수 있습니다.

오른쪽으로 k번 회전한 후 인덱스 i에 있는 요소는 원래 배열의 인덱스 (i - k) mod n에 있던 요소와 같습니다. 파이썬의 모듈로 연산은 음수도 자연스럽게 처리해 주므로, 아래와 같이 한 줄로 해결할 수 있습니다.

def solve(nums, k, i):
    return nums[(i - k) % len(nums)]

nums = [2,7,9,8,10]
k = 3
i = 2
print(solve(nums, k, i))  # 출력: 10

이 방법은 시간 복잡도가 O(1)로 훨씬 효율적이며, k가 배열의 길이보다 큰 경우에도 항상 정확하게 동작합니다.