Pythonでソート済みリストの各要素を2乗して昇順に並べたリストを取得する方法

ソート済みの数値リストが与えられたとき、各要素を2乗した結果を再びソートされた順序で出力することを考えます。入力には負の数や0が含まれる場合もある点に注意が必要です。2乗すると負の数は正の数になるため、単純に2乗するだけでは順序が崩れてしまいます。

例えば、入力が [-12, -6, -5, -2, 0, 1, 2, 4, 8, 9, 10, 15, 18, 20, 35, 38, 69] の場合、出力は次のようになります。

[0, 1, 4, 4, 16, 25, 36, 64, 81, 100, 144, 225, 324, 400, 1225, 1444, 4761]

解決手順

この問題は以下の手順で解くことができます。

• 結果を格納するための新しいリスト L を作成する
• 元のリスト nums の各要素 e について、e²(2乗した値)を L に追加する
• L をソートして返す

実装例

実際のコードは非常にシンプルです。Pythonの内包表記と組み込み関数 sorted() を使えば、わずか1行で処理できます。

class Solution:
    def solve(self, nums):
        return sorted(x * x for x in nums)

ob = Solution()
nums = [-12, -6, -5, -2, 0, 1, 2, 4, 8, 9, 10, 15, 18, 20, 35, 38, 69]
print(ob.solve(nums))

入力

[-12, -6, -5, -2, 0, 1, 2, 4, 8, 9, 10, 15, 18, 20, 35, 38, 69]

出力

[0, 1, 4, 4, 16, 25, 36, 64, 81, 100, 144, 225, 324, 400, 1225, 1444, 4761]

計算量について

この方法の計算量は、2乗の計算にO(n)、ソートにO(n log n)かかるため、全体で O(n log n) となります。

さらに効率的な方法(双方向ポインタ法)

入力リストがすでにソートされていることを利用すると、両端から比較する2ポインタ法(two-pointer technique)を使うことで、O(n)の時間で解くことも可能です。

• 絶対値が最も大きい要素ほど、2乗した値も大きくなる性質を利用します。
• リストの左端と右端を比較し、絶対値の大きい方の2乗を結果リストの末尾から埋めていきます。

def solve(nums):
    n = len(nums)
    result = [0] * n
    left, right = 0, n - 1
    for i in range(n - 1, -1, -1):
        if abs(nums[left]) > abs(nums[right]):
            result[i] = nums[left] ** 2
            left += 1
        else:
            result[i] = nums[right] ** 2
            right -= 1
    return result

この方法なら追加のソートが不要になり、大規模なデータセットでも高速に動作します。手軽さを重視するなら最初の sorted() を使う方法、パフォーマンスを重視するなら2ポインタ法を選ぶとよいでしょう。