Cyclic Rotation of Numbers in a List in Python

I am trying to do cyclic rotation of numbers in List in Python.

For example,

An array A consisting of N integers is given. Rotation of the array means that each element is shifted right by one index, and the last element of the array is moved to the first place. For example, the rotation of array A = [3, 8, 9, 7, 6] is [6, 3, 8, 9, 7] (elements are shifted right by one index and 6 is moved to the first place).

The goal is to rotate array A K times; that is, each element of A will be shifted to the right K times.

For example, given

A = [3, 8, 9, 7, 6]
K = 3

the function should return [9, 7, 6, 3, 8]. Three rotations were made:

[3, 8, 9, 7, 6] -> [6, 3, 8, 9, 7]
[6, 3, 8, 9, 7] -> [7, 6, 3, 8, 9]
[7, 6, 3, 8, 9] -> [9, 7, 6, 3, 8]

I have written a function that is supposed to do the above task. Here is my code:

def sol(A, K):
    for i in range(len(A)):
        if (i+K) < len(A):
            A[i+K] = A[i]
        else:
            A[i+K - len(A)] = A[i]
    return A

A = [3, 8, 9, 7, 6]
K = 3

# call the function
sol(A,K)
[9, 3, 8, 3, 8]

I am getting [9, 3, 8, 3, 8] instead of [9, 7, 6, 3, 8].

Can anyone help me with the above code ?

Solution

Let's take a look at what happens if K = 1, on just the first iteration:

def sol(A, K):
    for i in range(len(A)):  # i = 0
        if (i+K) < len(A):   # i + 1 < 5
            A[i+K] = A[i]    # A[1] = A[0]
        else:
            A[i+K - len(A)] = A[i]
        # A is now equal to [3, 3, 9, 7, 6] - the element at A[1] got overwritten
    return A

The problem is that you don't have anywhere to store the elements you'd be overwriting, and you're overwriting them before rotating them. The ideal solution is to create a new list, populating it with rotated elements from the previous list, and return that. If you need to modify the old list, you can copy elements over from the new list:

def sol(A, K):
    ret = []
    for i in range(len(A)):
        if (i + K) < len(A):
            ret.append(A[i + K])
        else:
            ret.append(A[i + K - len(A)])
    return ret

Or, more concisely (and probably how your instructor would prefer you solve it), using the modulo operator:

def sol(A, K):
    return [
       A[(i + K) % len(A)]
       for i in range(len(A))
    ]

Arguably the most pythonic solution, though, is to concatenate two list slices, moving the part of the list after index K to the front:

def sol(A, K):
    return A[K % len(A):] + A[:K % len(A)]