# The Lazy Squares

> A sequence that never fully reveals itself.

Canonical URL: <https://datadriven.io/problems/the_lazy_squares>

Domain: Python · Difficulty: easy · Seniority: L4

## Problem

Given positive integer n, yield (as a generator) 1^2, 2^2, ..., n^2 in order. The test harness collects them into a list.

## Worked solution and explanation

### Why this problem exists in real interviews

This tests **generator functions with yield**, the simplest possible generator pattern. It probes whether a candidate understands lazy evaluation and the difference between returning a list and yielding values one at a time.

---

### Break down the requirements

#### Step 1: Loop from 1 to n inclusive

Generate each integer in the range.

#### Step 2: Yield the square of each integer

Use `yield` instead of appending to a list.

---

### The solution

**Generator function yielding squares**

```python
def lazy_squares(n: int):
    for i in range(1, n + 1):
        square = i * i
        yield square
```

> **Time and Space Complexity**
>
> **Time:** O(n) to yield all values. O(1) per yield.
> 
> **Space:** O(1). The generator stores only its current state, not all values.

> **Interviewers Watch For**
>
> Whether you use `yield` vs building a list and returning it. The whole point of a generator is lazy, memory-efficient evaluation.

> **Common Pitfall**
>
> Using `return` instead of `yield`. A function with `return` produces one value; a generator with `yield` produces a sequence.

---

## Common follow-up questions

- What is the difference between a generator function and a generator expression? _(Tests `(x*x for x in range(1, n+1))` as the expression equivalent.)_
- How would you consume only the first 3 squares? _(Tests `itertools.islice` or manual `next()` calls.)_
- What happens if you iterate the generator twice? _(Tests that generators are exhausted after one pass; you need to recreate them.)_

## Related

- [All practice problems](https://datadriven.io/problems)
- [Mock interview mode](https://datadriven.io/interview/the_lazy_squares)
- [Python Interview Questions](https://datadriven.io/python-interview-questions)
- [Data Engineering Interview Prep Guide](https://datadriven.io/data-engineer-interview-prep)
- [Daily Challenge](https://datadriven.io/daily)

---

Source: DataDriven (https://datadriven.io). 100% free data engineering interview prep. Live code execution against Postgres 16, Python 3.11, and Spark sandboxes. No paywall, no premium tier, no signup gate.