LeetCode / Pascal's Triangle II

10482 ワード
Python leetcode Python テキストリンク

（ブログ記事からの転載）

AtCoderのBeginner Contentに初参戦してみました。

結果は緑(こちらのページによれば "ソフトウェアエンジニアとして申し分ない実力です" のレベル)相当には後一歩及ばずというところ。まだまだ精進です。

さて、今日の問題。

[https://leetcode.com/problems/pascals-triangle-ii/]

Given a non-negative index k where k ≤ 33, return the kth index row of the Pascal's triangle.

Note that the row index starts from 0.

In Pascal's triangle, each number is the sum of the two numbers directly above it.

Example:

Input: 3
Output: [1,3,3,1]

Follow up:

Could you optimize your algorithm to use only O(k) extra space?

今度はパスカルの三角形のk行目の数列を返すことが求められています。
前回記事の問題を解けたのであれば解くこと自体は造作もないですが、$O(k)$の省メモリな解法を考えるところで深みがあります。

解答・解説

解法1

前回記事にて以下のコードを示しましたが、

class Solution:
    def generate(self, numRows: int) -> List[List[int]]:
        ans = []
        if numRows > 0:
            ans.append([1])
            for _ in range(numRows-1):
                l_pre = [0]+ans[-1]+[0]
                l_nxt = [l_pre[i]+l_pre[i+1] for i in range(len(l_pre)-1)]
                ans.append(l_nxt)
        return ans

これを少し改変することでも解くことができました。

class Solution:
    def getRow(self, rowIndex: int) -> List[int]:
        l_pre = [1]
        if rowIndex > 0:
            for _ in range(rowIndex):
                l_pre = [0]+l_pre+[0]
                l_nxt = [l_pre[i]+l_pre[i+1] for i in range(len(l_pre)-1)]
                l_pre = l_nxt
        return l_pre
解法2

パスカルの三角形の各数列の計算は、以下のように上段の数列の片端に[0]を付けた二つの数列を足し合わせることでも計算できます。

こちらの方が解法1よりすっきりしていますね。

これを利用した解法が以下。

class Solution(object):
    def getRow(self, rowIndex):
        """
        :type rowIndex: int
        :rtype: List[int]
        """
        row = [1]
        for _ in range(rowIndex):
            row = [x + y for x, y in zip([0]+row, row+[0])]
        return row
解法3

最も空間計算量を抑えた解法はこちらかと思います。
解に必要な要素数のリストを用意し、その中でIterativeに数列を計算していく省メモリな解法です。

class Solution(object):
    def getRow(self, rowIndex):
        """
        :type rowIndex: int
        :rtype: List[int]
        """
        ans = [0]*(rowIndex+1)
        ans[0] = 1
        for i in range(1,rowIndex+1):
            for j in range(i,0,-1):
                ans[j] += ans[j-1]
        return ans

具体的に出力変数ansの変化を見てみると、イメージが湧くと思います。(rowIndex=5)

[1, 1, 0, 0, 0, 0]
[1, 2, 1, 0, 0, 0]
[1, 3, 3, 1, 0, 0]
[1, 4, 6, 4, 1, 0]
[1, 5, 10, 10, 5, 1]