Coding - Nested Structure Flatten and Unflatten - xAI Interview Question | DarkInterview

Nested Structure Flatten and Unflatten

CodingPhoneSoftware EngineerReported Sep, 2026

Problem Overview

Many production pipelines need to transform hierarchical data into a linear representation and then reconstruct the original shape later. This shows up in model input preprocessing, feature extraction, serialization layers, and tree-structured configuration systems.

In this problem, a nested Python structure may contain:

list
tuple
dict

All leaf values are guaranteed to be integers. [Source: darkinterview.com]

Your tasks are:

Flatten the nested structure into a single list of integers
Unflatten a flat list of integers back into the same nested shape as a template

Traversal Rules

To make the output deterministic, use the following traversal order:

For a list, visit elements from left to right
For a tuple, visit elements from left to right
For a dict, visit values in the dictionary's iteration order

Assume standard Python 3.7+ dictionary ordering semantics. [Source: darkinterview.com]

Example

structure = {
    "a": [1, 2, 3],
    "b": {
        "c": [{"d": 4}],
        "e": 5,
    },
}

# Flatten
flatten(structure)  # [1, 2, 3, 4, 5]

# Unflatten using the same shape
unflatten([6, 7, 8, 9, 0], structure)
# {
#     "a": [6, 7, 8],
#     "b": {
#         "c": [{"d": 9}],
#         "e": 0,
#     },
# }

Part 1: Flatten a Nested Structure

Problem Statement

Write a function that takes a nested structure and returns a flat list containing every integer leaf in traversal order.

from typing import Any

def flatten(structure: Any) -> list[int]:
    """
    Flatten a nested structure of lists, tuples, dicts, and integer leaves.

    Args:
        structure: A nested Python object whose internal nodes are list, tuple,
            or dict, and whose leaves are integers.

    Returns:
        A flat list of integers in traversal order.
    """
    pass

Test Cases

# Test 1: Nested dict + list
structure = {"a": [1, 2, 3], "b": {"c": [{"d": 4}], "e": 5}}
assert flatten(structure) == [1, 2, 3, 4, 5]

# Test 2: Tuple and list mixed together
structure = (1, [2, (3, 4)], {"x": 5})
assert flatten(structure) == [1, 2, 3, 4, 5]

# Test 3: Single integer leaf
structure = 42
assert flatten(structure) == [42]

# Test 4: Empty containers
structure = {"a": [], "b": (), "c": {}}
assert flatten(structure) == []

# Test 5: Deep nesting
structure = [[[[1]]], {"a": {"b": {"c": 2}}}, (3,)]
assert flatten(structure) == [1, 2, 3]

# Test 6: Dict order matters
structure = {"first": 1, "second": [2, 3], "third": 4}
assert flatten(structure) == [1, 2, 3, ]

Solution: Recursive DFS for Flatten

The simplest solution is a depth-first traversal:

If the current value is an integer, append it to the output
If it is a list or tuple, recurse over each element in order
If it is a dict, recurse over each value in iteration order

from typing import Any

def flatten(structure: Any) -> list[int]:
    result: list[int] = []

    def dfs(node: Any) -> None:
        if isinstance(node, int):
            result.append(node)
        elif isinstance(node, list) or isinstance(node, tuple):
            for item in node:
                dfs(item)
        elif isinstance(node, dict):
            for value in node.values():
                dfs(value)
        else:
            raise TypeError(f"Unsupported node type: {type(node)}")

    dfs(structure)
    return result

Complexity: [Source: darkinterview.com]

Time: O(N), where N is the total number of nodes/leaves visited
Space: O(L + H), where L is the output size and H is recursion depth

Part 2: Unflatten Back Into a Template

Problem Statement

Now implement the inverse transformation. Given:

flat_list: a list of integers
structure: a nested object used as a template

Create a new object with the exact same shape and container types as structure, but replace each integer leaf with the next value from flat_list. Dictionary keys should be preserved exactly; only leaf values are replaced.

If flat_list contains too few or too many integers for the template, raise a ValueError. [Source: darkinterview.com]

from typing import Any

def unflatten(flat_list: list[int], structure: Any) -> Any:
    """
    Rebuild a nested object with the same shape as `structure`.

    Args:
        flat_list: The replacement integer values, consumed in order.
        structure: Template object whose shape determines the output.

    Returns:
        A new nested object with the same structure and container types.

    Raises:
        ValueError: If flat_list length does not exactly match the number of
            integer leaves required by structure.
    """
    pass

Example

template = {
    "a": [1, 2, 3],
    "b": {
        "c": [{"d": 4}],
        "e": 5,
    },
}

result = unflatten([6, 7, 8, 9, 0], template)

assert result == {
    "a": [6, 7, 8],
    "b": {
        "c": [{"d": 9}],
        "e": 0,
    },
}

Test Cases

# Test 1: Basic reconstruction
template = {"a": [1, 2, 3], "b": {"c": [{"d": 4}], "e": 5}}
expected = {"a": [6, 7, 8], "b": {"c": [{"d": 9}], "e": 0}}
assert unflatten([6, 7, 8, 9, 0], template) == expected

# Test 2: Preserve tuple type
template = (1, [2, 3], {"x": 4})
result = unflatten([9, 8, 7, 6], template)
assert result == (9, [8, 7], {"x": 6})
assert isinstance(result, tuple)

# Test 3: Empty containers remain empty
template = {"a": [], "b": (), "c": {}}
assert unflatten([], template) == {"a": [], "b": (), "c": {}}

# Test 4: Single leaf
template = 100
assert unflatten([5], template) == 5

# Test 5: Too few values
template = [1, 2, ]
:
    unflatten([, ], template)
     , 
 ValueError:
    


template = {: }
:
    unflatten([, ], template)
     , 
 ValueError:
    


template = {: [], : (), : {}}
:
    unflatten([], template)
     , 
 ValueError:

Solution: Recursive Reconstruction with an Iterator

For unflatten, the key idea is to walk the template in the exact same order and consume values from flat_list one by one.

Two common implementations work well:

Maintain a mutable index into flat_list
Convert flat_list into an iterator and call next(...) at each leaf

The iterator version is especially clean because container-building logic stays local to each recursive call. [Source: darkinterview.com]

from typing import Any

def unflatten(flat_list: list[int], structure: Any) -> Any:
    values = iter(flat_list)

    def build(node: Any) -> Any:
        if isinstance(node, int):
            try:
                return next(values)
            except StopIteration as exc:
                raise ValueError("flat_list has fewer values than required") from exc
        if isinstance(node, list):
            return [build(item) for item in node]
        if isinstance(node, tuple):
            return tuple(build(item) for item in node)
        if isinstance(node, dict):
            return {key: build(value) for key, value in node.items()}
        raise TypeError(f"Unsupported node type: {type(node)}")

    rebuilt = build(structure)

    try:
        next(values)
        raise ValueError("flat_list has more values than required")
    except StopIteration:
        return rebuilt

Complexity:

Time: O(N)
Space: O(H) auxiliary recursion space, not counting the output

Key Insight

flatten and unflatten are two sides of the same traversal contract. The hardest part is not the recursion itself, but preserving a deterministic traversal order across lists, tuples, and dictionaries. If that order is inconsistent, reconstruction breaks even when the code looks otherwise correct.

Additional Discussion

A good follow-up is to extract a shared traversal helper so both functions rely on exactly the same ordering rules.
Another strong extension is replacing integers with arbitrary leaves and passing a predicate like is_leaf(node).
For very deep inputs, an explicit stack can avoid Python recursion depth limits.