競技プログラミング Python

2Sum and 3Sum - Leetcode 1 & 15

Python: Competitive Programming 2/N

公開日: 2023-09-29
更新日: 2023-09-29

  Table of Contents

Problem 1: 2Sum

intgerを格納したリスト numsと integer targetが与えられたとき, numsの要素を２つ抽出したとき, その合計がtargetとなるようなnumsのindexをlistで出力してください.

  • numsのうち同じindexのもの２つは抽出できない
  • 解はユニークに必ず存在するとする

例えば,

  • Input: nums = [2,7,11,15], target = 9
  • Output: [0,1]

nums[0] + nums[1] == targetとなるため上記の出力となる.

また, 以下のようにnumsの要素は重複している可能性もあるが, 必ずユニークな答えが存在するとする.

  • Input: nums = [3,3], target = 6
  • Output: [0,1]

$O(N^2)$ Solution: Nested Loop

1
2
3
4
5
6

class Solution:
    def twoSum(self, nums: List[int], target: int) -> List[int]:
        for left_idx in range(0, len(nums)):
            for right_idx in range(left_idx+1, len(nums)):
                if nums[left_idx] + nums[right_idx] == target:
                    return [left_idx, right_idx]

この答えは二重ループで, 最悪の場合if文がnums sizeを$N$としたとき,

\[\frac{N(N+1)}{2} = \frac{1}{2}(N^2 + N)\]

呼ばれるので, Time Complexityは$O(N^2)$となっていることがわかる.

$O(N)$ Solution: Hash Map

1
2
3
4
5
6
7
8

class Solution:
    def twoSum(self, nums: List[int], target: int) -> List[int]:
        thresh = target - min(nums)
        d = {j:i for i, j in enumerate(nums) if j <= thresh}
        for i, num in enumerate(nums):
            diff = target - num
            if diff in d and d[diff] != i:
                return [i, d[diff]]

これは仮定の一つの「解はユニークに必ず存在するとする」に依拠した解答となっています.

1

d = {j:i for i, j in enumerate(nums) if j <= thresh}

でnumsの要素をkey, indexをvalueとした（重複要素の場合は大きい方のindexを用いた）dict(Hash Map)を $O(N)$の計算量で定義します. このとき, target - min(target)より大きい要素を格納してもメモリスペースの無駄なので, 予め除去しておきます.

その後, numsの各要素numに対して,

1

diff = target - num

を計算し, diffをkeyとして持つitemが存在するか, $O(N)$の計算量で判定しています. そのため, この解はlinear time comlexityの解答となっています.

Column: 計算量

計算量とはアルゴリズムの効率性を表す指標で, 時間計算量と空間計算量の２つがあります.

  • 時間計算量, Time Complexity: 入力データに対してアルゴリズムを実行するための必要な時間の指標
  • 空間計算量, Memory Space: 計算に必要なメモリ量の指標

$O(N \log N)$ Solution: sorted array

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

from typing import List

class Solution:
    def twoSum(self, nums: List[int], target: int) -> List[int]:
        sorted_num = sorted(nums)
        left_idx = 0
        right_idx = len(nums) - 1

        while left_idx < right_idx and right_idx < len(nums):
            left_val, right_val = sorted_num[left_idx], sorted_num[right_idx]
            tmp = left_val + right_val
            if tmp < target:
                left_idx += 1
            elif tmp > target:
                right_idx -= 1
            else:
                first_idx = nums.index(left_val)
                second_idx = len(nums) - 1 - nums[::-1].index(right_val)
                return [first_idx, second_idx]

この解のtime complexity, $O(N\log(N))$ は以下のように分解して考えることができます:

  • sorted(nums): worst caseにおいて $O(N \log(N))$
  • while loop: sorted arrayの要素は一回ずつしか読み込まれないので $O(N)$

なお, Leetcodeで手法別計算量を確認してみると, Hash Map/sorted array Solutionの優秀さが以下のように確認できます.

手法 Runtime Memory
Nested Loop 2467 ms 17 MB
Hash Map 73 ms 17.3 MB
sorted array 68 ms 17 MB

Problem 15: 3Sum

intgerを格納したリスト numsが与えられたとする. indexが異なるという条件のもと, numsの要素を3つ抽出したとき, その合計がtargetとなるようなnumsのlistで出力してください.

  • ３つの要素の組について重複するものは除去した上でnumsのlistを出力

例として,

  • Input: nums = [-1,0,1,2,-1,-4]
  • Output: [[-1,-1,2],[-1,0,1]]

このとき, [[-1,-1,2],[-1,0,1],[-1, 2, -1]]のように出力すると, [-1, 2, -1]と[-1,-1,2]は組み合わせが 重複しているので片方の出力は認められない.

Solution

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48

from typing import List


class Solution:
    """
    Problem
        https://leetcode.com/problems/3sum/

        Given an integer array nums, return all the triplets [nums[i], nums[j], nums[k]] 
        such that i != j, i != k, and j != k, and nums[i] + nums[j] + nums[k] == 0.
        Notice that the solution set must not contain duplicate triplets.
    
    Related
        https://leetcode.com/problems/two-sum/
    
    Note
        this time, required to retund list of values, not index
    """
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        nums.sort()
        res = []
        max_idx = len(nums)-1
        prev = None
        for first_idx in range(max_idx):
            # reduce the search area
            if nums[first_idx] > 0:
                return res
            if nums[first_idx] == prev:
                continue
            
            # Main Part
            target = 0 - nums[first_idx]
            second_idx = first_idx + 1
            third_idx = max_idx
            while second_idx < third_idx and third_idx <= max_idx:
                tmp = nums[second_idx] + nums[third_idx]
                if tmp > target:
                    third_idx -= 1
                elif tmp < target:
                    second_idx += 1
                else:
                    res.append([nums[first_idx], nums[second_idx], nums[third_idx]])
                    second_idx += 1
                    while nums[second_idx] == nums[second_idx - 1] and second_idx < third_idx:
                        second_idx += 1
            prev = nums[first_idx]
        
        return res
  • Runtime complexyは$O(N^2)$
  • Runtime: 739 ms
  • Memory: 20.3 MB

References

Share Buttons
 Share on:
Feature Tags
Leave a Comment
(注意：GitHub Accountが必要となります）
CONTENTS