【競技プログラミング】「Merge Sorted Array」with Python | LeetCode #88

Python

こちらは「Merge Sorted Array」のPythonによる解答例となっております。この問題は、2つのソートされた整数配列が与えられたとき、2つ目の配列を最初の配列にマージし、結果をソートされた状態で返す問題です。

問題

2つの整数配列 nums1 と nums2 がソートされていると仮定します。nums1 には十分な空きがあり、nums2 の要素を格納するための空きがあります(nums1 と nums2 の合計要素数は m と n です)。nums2 の要素を nums1 にマージして、マージした結果をソートされた状態で返します。

  • 0 <= m, n <= 200
  • 1 <= m + n <= 200
  • nums1.length == m + n
  • nums2.length == n
  • 10^9 <= nums1[i], nums2[i] <= 10^9

解法

この問題は、2つの配列をマージするアルゴリズムが必要です。2つの配列がソートされているため、マージする際には2つの配列の要素を比較して小さい方から新しい配列に追加することができます。具体的なアルゴリズムは以下の通りです。

  1. nums1nums2 の末尾から比較し、大きい方を nums1 の末尾に格納します。
  2. nums1nums2 のどちらかが終了するまで、1を繰り返します。
  3. nums2 が終了した場合、nums1 の開始位置から終了位置までをソートします。

Pythonコード

以下のコードが、「Merge Sorted Array」を解くためのPythonプログラムとなります。

### Function ##################################################################

def solve(num=0):
    
    solver  = Solution().merge
    
    if 1 <= num <= Input().input: id_LIST = [num]
    else: id_LIST = [id+1 for id in range(Input().input)]
    
    for id in id_LIST:
        data    = Input(id)
        ans     = solver(data.nums1, data.m, data.nums2, data.n)
        result  = "OK" if ans==data.ans else "NG"
        print(f"[{str(id).zfill(3)}] {result}")

### Class #####################################################################

class Solution(object):
    
    def merge(self, nums1, m, nums2, n):
        """
        :type nums1: List[int]
        :type m: int
        :type nums2: List[int]
        :type n: int
        :rtype: None Do not return anything, modify nums1 in-place instead.
        """
        
        while m > 0 and n > 0:
            
            if nums1[m - 1] > nums2[n - 1]:
                nums1[m + n - 1] = nums1[m - 1]
                m -= 1
            
            else:
                nums1[m + n - 1] = nums2[n - 1]
                n -= 1
        
        nums1[:n] = nums2[:n]
        
        return nums1
    
    # def merge(self, nums1, m, nums2, n):
    #     """
    #     :type nums1: List[int]
    #     :type m: int
    #     :type nums2: List[int]
    #     :type n: int
    #     :rtype: None Do not return anything, modify nums1 in-place instead.
    #     """
    #     for id in range(len(nums1)-m): nums1.pop()
    #     nums1.extend(nums2)
    #     nums1.sort()

#==============================================================================

class Input:
    
    def __init__(self, num=0):
        
        if num == 1:
            self.nums1 = [1,2,3,0,0,0]
            self.nums2 = [2,5,6]
            self.m = 3
            self.n = 3
            self.ans    = [1,2,2,3,5,6]
        
        elif num == 2:
            self.nums1 = [1]
            self.nums2 = []
            self.m = 1
            self.n = 0
            self.ans    = [1]

        elif num == 3:
            self.nums1 = [0]
            self.nums2 = [1]
            self.m = 0
            self.n = 1
            self.ans    = [1]
        
        else:
            self.input  = 3

### Main ######################################################################

if __name__ == "__main__":

    solve()

###############################################################################