HTMLify

vql05pvdcr.py
Views: 22 | Author: Unknown
def flatten_and_count(matrix) -> tuple[list[int], int]:
    ic = 0
    fl = []
    for line in matrix:
        if type(line) == list:
            fl_, ic_ = flatten_and_count(line)
            fl.extend(fl_)
            ic += ic_
        else:
            fl.append(line)
        ic += 1
    return fl, ic


# Tests

Tests = [
    {
        "matrix": [[1, 2, 3, [22, [1111, 2222, 333], 33, 44]], [4, 5, 6], [7, 8, 9]],
        "expected": ([1, 2, 3, 22, 1111, 2222, 333, 33, 44, 4, 5, 6, 7, 8, 9], 20),
    },
    {
        "matrix": [[1, 2], [3, 4], [5]],
        "expected": ([1, 2, 3, 4, 5], 8),
    }
]

for test in Tests:
    out = flatten_and_count(test["matrix"])
    if out != test["expected"]:
        print("test faild")
        print("input:\n", test["matrix"])
        print("expected:\n", test["expected"])
        print("got:\n", out)
    else:
        print("test pass")

Comments