[Python]斜めに情報を引き継ぐ

概要

ABC183 E - Queen on Grid で、右下へ移動する方法の数の累積和を記録する必要があった。
直前と現在の 2 つの配列を用意し、直前の値を参照して現在の値を更新する方法もあるが、1 つの配列を使って書くこともできる。

2 つの配列で引き継ぐ

右上の値を 10 倍するサンプル。

import numpy as np

h, w = 3, 4
A = np.zeros((h, w), int)
prev = [0, 1, 2, 3]

for y in range(h):
    A[y] = prev
    cur = [0]*w
    for x in range(w):
        cur[x] = w+y if x == w-1 else 10*prev[x+1]
    prev = cur

print(A)

# 結果
# [[  0   1   2   3]
#  [ 10  20  30   4]
#  [200 300  40   5]]

配列の境界や初期条件を与える実装がやや複雑になる。

1 つの配列で引き継ぐ

先ほどと同じ、右上の値を 10 倍するサンプル。

import numpy as np

h, w = 3, 4
A = np.zeros((h, w), int)
cur = [i for i in range(w+h-1)]

for y in range(h):
    for x in range(w):
        k = x+y
        A[y, x] = cur[k]
        cur[k] *= 10 
print(A)

# 結果
# [[  0   1   2   3]
#  [ 10  20  30   4]
#  [200 300  40   5]]

少し慣れが必要だが、配列の境界を意識せずに書ける。また、初期条件を一か所に定義できることにも利点がある。

同様に、左上の値を 10 倍するサンプル。

import numpy as np

h, w = 3, 4
A = np.zeros((h, w), int)
cur = [i for i in range(w+h-1)]

for y in range(h):
    for x in range(w):
        k = x-w-y
        A[y, x] = cur[k]
        cur[k] *= 10

print(A)

# 結果
# [[  2   3   4   5]
#  [  1  20  30  40]
#  [  0  10 200 300]]