リストを使わずyieldを使って数列を生成したい

ある画像処理用のアルゴリズムのPythonによる実装です。目的の数列は得られるのですが、リストを使っているためにメモリ効率が悪いです。Pythonにはジェネレータがあることを知り、リストを使わずyieldを使ってスマートにメモリ効率の良いコードが書けるのではないかと思ったのですが、Pythonを学んで日が浅く、わからずにいます。

Python
1import sys
2
3
4def create(depth, left, right):
5    if left > right:
6        return None
7    else:
8        return bin_node(depth + 1, left, right)
9
10
11def bin_node(depth, left, right):
12    center = (left + right) // 2
13    return {
14            'depth': depth,
15            'center': center,
16            'leftNode': create(depth, left, center - 1),
17            'rightNode': create(depth, center + 1, right)
18            }
19
20
21def get_max_depth(node):
22    if node is None:
23        return 0
24    else:
25        return max([node['depth'], get_max_depth(node['leftNode']), get_max_depth(node['rightNode'])])
26
27
28def traverse_in_depth(node, depth):
29    list = []
30    if node is not None and node['depth'] == depth:
31        list.append(node['center'])
32    elif node is not None and node['depth'] < depth:
33        list.extend(traverse_in_depth(node['leftNode'], depth))
34        list.extend(traverse_in_depth(node['rightNode'], depth))
35    return list
36
37
38def bin_node_iterator(left, right):
39    root_node = create(0, left, right)
40
41    max_depth = get_max_depth(root_node)
42
43    list = []
44    for depth in range(1, max_depth + 1):
45        list.extend(traverse_in_depth(root_node, depth))
46
47    return list
48
49
50left = int(sys.argv[1])
51right = int(sys.argv[2])
52print(left, right)
53for n in bin_node_iterator(left, right):
54    print(n, end=" ")

試したこと
（traverse_in_depth関数、bin_node_iterator関数以外は同じ）

Python
1def traverse_in_depth(node, depth):
2    if node is not None and node['depth'] == depth:
3        yield node['center']
4    elif node is not None and node['depth'] < depth:
5        yield traverse_in_depth(node['leftNode'], depth)
6        yield traverse_in_depth(node['rightNode'], depth)
7
8
9def bin_node_iterator(left, right):
10    root_node = create(0, left, right)
11
12    max_depth = get_max_depth(root_node)
13
14    for depth in range(1, max_depth + 1):
15        yield traverse_in_depth(root_node, depth)

行動規範の内容に同意します

回答1件

ベストアンサー

◯ 提案

全体のコードを読めていないので、直るかはわからないのですが
traverse_in_depth 関数を呼び出す際に yield from を使ってみると、どうでしょうか？

# 修正前
def traverse_in_depth(node, depth):
    if node is not None and node['depth'] == depth:
        yield node['center']
    elif node is not None and node['depth'] < depth:
        yield traverse_in_depth(node['leftNode'], depth)
        yield traverse_in_depth(node['rightNode'], depth)


def bin_node_iterator(left, right):
    root_node = create(0, left, right)

    max_depth = get_max_depth(root_node)

    for depth in range(1, max_depth + 1):
        yield traverse_in_depth(root_node, depth)

# 修正後
def traverse_in_depth(node, depth):
    if node is not None and node['depth'] == depth:
        yield node['center']
    elif node is not None and node['depth'] < depth:
        yield from traverse_in_depth(node['leftNode'], depth)  # <--- yield from を使います。
        yield from traverse_in_depth(node['rightNode'], depth)  # <--- yield from を使います。


def bin_node_iterator(left, right):
    root_node = create(0, left, right)

    max_depth = get_max_depth(root_node)

    for depth in range(1, max_depth + 1):
        yield from traverse_in_depth(root_node, depth)  # <--- yield from を使います。

◯ 補足

yield を使った関数を「ジェネレータ関数」と言います。
「ジェネレータ関数」は「ジェネレータイテレータ」というオブジェクトを返します。

traverse_in_depth 関数は「ジェネレータイテレータ」を返しています。
修正前の bin_node_iterator 関数は traverse_in_depth 関数が返した
「ジェネレータイテレータ」を１つずつ返してしまっています。

python
1# コピペで動きます。
2def f():
3    yield g1()
4    yield g2()
5
6def g1():
7    yield 0
8    yield 1
9
10def g2():
11    yield 2
12    yield 3
13
14[i for i in f()]
15# [<generator object g1 at 0x100e9db48>, <generator object g2 at 0x100e9dba0>]

今回は「ジェネレータイテレータ」の中身が欲しいので、
yield from を使います。

python
1# コピペで動きます。
2def f():
3    yield from g1()
4    yield from g2()
5
6def g1():
7    yield 0
8    yield 1
9
10def g2():
11    yield 2
12    yield 3
13
14[i recursive.pyfor i in f()]
15# [0, 1, 2, 3]