はじめに
競プロ中、2次元配列の変化を観察したい時があります。
DPとか、DPとか。
1次元配列ならまだしも、2次元配列を逐次print
するとかなり見にくいので、
2次元配列の変化をじーっと眺めるための関数を用意しました。
完成品はこちら↓
from time import sleep
def print_2d_overwrite(list2d, val_width=5, sleep_sec=1.0, header=None, highlights=[]):
res_str = f'{header}\n' if header is not None else ''
for i, row in enumerate(list2d):
for j, v in enumerate(row):
v_str = str(v).rjust(val_width)
if (i,j) in highlights:
v_str = f'\033[31m{v_str}\033[0m'
res_str += f' {v_str}'
res_str += '\n'
new_line_cnt = res_str.count('\n')
res_str += f'\033[{new_line_cnt}A'
print(res_str, end='')
sleep(sleep_sec)
# --- sample ---
n = 3
list2d = [ [0]*n for _ in range(n)]
for i in range(n):
for j in range(n):
list2d[i][j] = 3**(i+j)
print_2d_overwrite(list2d, val_width=3, header=(i,j), highlights=[(i,j)])
ポイントは以下です。
- 標準出力を上書きする
- 値の長さ(表示桁)を揃える
- 文字列の色を変える
- sleepする(
time.sleep()
するだけ)
(2020/11/19 追記)
コメントで、テクニカルな書き方をいただきました。
標準出力を上書きする
関連する記事 が昨日投稿されていてビックリ。
自分が実装の参考にした記事は こちら です。
これらの記事で解説は十分なので、今回は要点だけ説明。
\033[nA
で、カーソルをn行上に移動できます。
表示したい2次元配列を文字列 (res_str)
に変換し、その中の 改行\n
の個数を数え、文字列の最後に追加しています。f文字列を使うと楽。
print
でend=''
を忘れずに。
new_line_cnt = res_str.count('\n')
res_str += f'\033[{new_line_cnt}A'
print(res_str, end='')
値の長さ(表示桁)を揃える
こんな感じで長さが揃ってないと見にくい
3 5 345
235 -48 123
1 2 3
ので、各要素の長さ(表示桁)を揃えたいです。
3 5 345
235 -48 123
1 2 3
今回は、str.rjust(n)
を利用してこれを実現します。
文字列を右寄せにし、n文字に足りない分は第2引数で指定した文字で埋めます(何も指定しない場合はスペース埋め)。
v_str = str(v).rjust(val_width)
文字列の色を変更する
\033[31m色変したい文字列\033[0m
でOKです。
30m
~37m
で色設定可能。31m
は赤です。
(詳細は こちら)
関数内で、highlights
に指定された値の色を変更しています。こちらもf文字列を使うと楽。
if (i,j) in highlights:
v_str = f'\033[31m{v_str}\033[0m'
おわりに
もしかしたらデバッグで使えるかもしれないので、競プロ用スニペットに加えておきました。
sleep
の部分をキー入力にしたらもっと便利かもと思ったのですが、Macだとsudo
しないとkeyboard
が使えず面倒だったので、このままで。