DSA Fruit into Baskets

The problem

Fruit into Baskets finds the longest stretch of a row of trees where at most two kinds of fruit appear. Each of the two baskets holds one kind, and a stretch is picked in one unbroken run.

The answer is the length of that stretch. The row may be empty. A stretch with one kind is always allowed, because one kind fits the two-kind budget.

cases = [
    (["apple", "pear", "apple", "pear", "plum", "pear"], 4),
    (["plum", "plum", "apple", "plum"], 4),
    (["apple", "pear", "plum"], 2),
    (["apple", "apple"], 2),
]
for row, expected in cases:
    print("row", row, "longest:", expected)

Read one row at a time. The first row answers 4, from its first four trees. The third row holds three kinds in three trees, so no stretch of length 3 is allowed.

Intuition: starting at the first tree is a trap

The first idea starts at the first tree and picks until a third kind appears. That gives one valid stretch. It says nothing about later stretches.

The trap is real in the row below. The prefix stops after two trees. A later stretch of four trees fits the same two baskets.

row = ["plum", "apple", "pear", "pear", "pear"]
kinds = []
prefix = 0
for fruit in row:
    if fruit not in kinds:
        if len(kinds) == 2:
            break
        kinds.append(fruit)
    prefix += 1
print("prefix length from the first tree:", prefix)
print("later stretch:", row[1:5], "length:", len(row[1:5]))

Starting position decides the prefix answer. The best stretch can begin anywhere in the row. Any method that only tries the start throws away most of the row.

Brute force: try every start

The brute force tries each tree as the start of a stretch. From a start it walks right until a third kind appears. A counter records how many trees get examined.

A kind is a distinct fruit name in the stretch. The walk from one start stops at the first tree that would make three kinds. That tree is not counted in the stretch.

def longest_brute(row):
    best = 0
    examined = 0
    for start in range(len(row)):
        kinds = set()
        for end in range(start, len(row)):
            examined += 1
            kinds.add(row[end])
            if len(kinds) > 2:
                break
            best = max(best, end - start + 1)
    return best, examined

best, examined = longest_brute(["apple", "pear", "apple", "pear", "plum", "pear"])
print("longest:", best)
print("trees examined:", examined)
best, examined = longest_brute(["plum", "plum", "apple", "plum"])
print("longest:", best)
print("trees examined:", examined)

The first row examined 18 tree visits for 6 trees. Each start repays the cost of walking over trees that earlier starts already covered. That repeat work is what the window removes.

The optimal solution: a window with a two-kind budget

A sliding window holds a count per kind, grows on the right, and shrinks from the left while it holds three kinds. That one sentence is the whole method.

A kind leaves the budget only when its count reaches zero. Dropping the count to zero and deleting the entry are one event. While two or fewer kinds remain, the window is valid and its length can update the best.

def longest_pick(row):
    counts = {}
    left = 0
    best = 0
    for right, fruit in enumerate(row):
        counts[fruit] = counts.get(fruit, 0) + 1
        while len(counts) > 2:
            counts[row[left]] -= 1
            if counts[row[left]] == 0:
                del counts[row[left]]
            left += 1
        best = max(best, right - left + 1)
    return best

cases = [
    (["apple", "pear", "apple", "pear", "plum", "pear"], 4),
    (["plum", "plum", "apple", "plum"], 4),
    (["apple", "pear", "plum"], 2),
    (["apple", "apple"], 2),
]
for row, expected in cases:
    result = longest_pick(row)
    print("row", row, "->", result, "expected:", expected, "match:", result == expected)

All four cases from the start of the article pass. The whole-row row passes because it never holds a third kind. Its window never shrinks at all.

Walkthrough: the six-tree row

The walk grows the window one tree at a time. The state is the window itself and its kind counts. The plum at index 4 forces the left edge forward until one kind drops out.

row = ["apple", "pear", "apple", "pear", "plum", "pear"]
counts = {}
left = 0
best = 0
for right, fruit in enumerate(row):
    counts[fruit] = counts.get(fruit, 0) + 1
    while len(counts) > 2:
        counts[row[left]] -= 1
        if counts[row[left]] == 0:
            del counts[row[left]]
        left += 1
    best = max(best, right - left + 1)
    print("window", row[left:right + 1], "- kinds:", sorted(counts), "- best:", best)

The window reaches four trees before the plum arrives. After the shrink it holds two trees, and it rebuilds to only three by the end of the row. The best value of 4 is read straight off the printed steps.

Complexity

The brute force walks right from every start, so a row of n trees can cost O(n squared) tree visits. The sliding window moves each edge across the row at most once, so its time is O(n).

The brute force stores the kinds of one stretch, at most three names before it stops, so its extra space is O(1). The window stores one count per kind inside it, and it never holds more than three kinds, so its extra space is also O(1).

for n in [6, 12, 24]:
    print("trees", n, "- brute visits in the worst case:", n * (n + 1) // 2, " window edge moves at most:", 2 * n)

Counts beat clocks here. At 24 trees the worst-case recount visits 300 positions. The window's two edges together move at most 48 steps.

Edge cases and pitfalls

Forgetting to delete a kind at count zero is the classic bug. The kind then still occupies a basket on paper. The window shrinks far more than it should, and valid stretches get cut short.

An empty row answers 0. A row of one kind answers its full length. Both fall out of the same code with no special case.

def longest_pick(row):
    counts = {}
    left = 0
    best = 0
    for right, fruit in enumerate(row):
        counts[fruit] = counts.get(fruit, 0) + 1
        while len(counts) > 2:
            counts[row[left]] -= 1
            if counts[row[left]] == 0:
                del counts[row[left]]
            left += 1
        best = max(best, right - left + 1)
    return best

print("empty row:", longest_pick([]))
print("one tree:", longest_pick(["plum"]))
print("one kind only:", longest_pick(["pear", "pear", "pear"]))
print("alternating two kinds:", longest_pick(["apple", "pear", "apple", "pear"]))
print("third kind at the very end:", longest_pick(["apple", "apple", "pear", "plum"]))

The alternating row never triggers a shrink. Two kinds can repeat forever inside the budget. Only a third distinct kind, not a long row, forces the left edge to move.