CodilityのCyclicRotationをPythonで解いてみた

こんにちは😄えびならいふです。
今回は前回に続きプログラミングに関する話題で、CodilityというプログラミングテストのようなもののLessonをPythonで解いてみた話。
今日はその中の最初のLesson 2 CyclicRotationという問題を解いてみた。

CyclicRotation

まずはCyclicRotationがどういう問題かの説明。 入力は配列Aと自然数Kの2つで、Aの中の要素をK回ローテーションしたときにどんな配列になるかを当てれば良い。
例えば、以下のようなAとKが入力として与えられた場合、

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

Aを1回ローテーションするとA=[6, 3, 8, 9, 7]となり、2回目はA=[7, 6, 3, 8, 9]、3回目はA=[9, 7, 6, 3, 8]となり、この例だと答えは[9, 7, 6, 3, 8]となる。
この問題は親切にヒントも与えられていて、

A = [0, 0, 0]
K = 1

の場合は答えが[0, 0, 0]となり、

A = [1, 2, 3, 4]
K = 4

の場合は答えが[1, 2, 3, 4]となると書かれている。
以上のように問題は非常にわかりやすくて単純明快なので次からは実際に解いてみる。

解いてみる

こういうプログラミングテストで注意しないといけないのは入力となる配列や変数の範囲。
配列の長さが0の時とかに処理に失敗してエラーが出るなんてこともあり得るので、最初に長さ0の場合の処理を書いておく。

if not A:
    return A

また、Aの要素が全部同じ場合はローテーションする必要がないのでこちらも書いてしまっておく*1。

if set(A) == 1:
    return A

配列Aの長さ分ローテーションすると元の配列に戻るため、実際には律儀にK回ローテーションする必要はなく、Kを配列Aの長さで割った余りの分だけローテーションすれば良いのでKを以下のようにしておく。
Kが0もしくは配列Aと同じ長さの場合はそのままAを返す。

len_A = len(A)
K = K % len(A)
if K == len_A or K == 0:
    return A

ここまでくればほぼ終わり。
for文で回して一つずつ動かしてもいけるかもしれないけど、配列の長さが大きくなると少し大変なので少しだけ工夫してみる。
といっても大した工夫ではなくて、単純にsliceを使う方法。
ローテーションと言っても、後ろの数字列を前に持っていってるに過ぎないので、配列をsliceで前後に分けてそれを入れ替えればOK。

n = len_A - K # 分割位置
# 前と後ろに分割
A_a = A[:n]
A_b = A[n:]
# 後ろの数字列の後ろに元々前にあったやつをくっつける
A_b.extend(A_a)
return A_b

終わりに

今回の実装でスコアはこんな感じになる。

*1:書かなくてもエラーになったりはしないが速度考えると書いておいた方が早い