Treasure and Cities Problem Using Dynamic Programming

The Treasure and Cities problem asks for the maximum profit from selectively visiting cities in order, where each visit pays a * t[i] for a colour match with the previous city or b * t[i] otherwise.

Greedy selection fails here. The algorithm is a variant of 0/1 knapsack and runs in O(n²) time.

Problem Statement

You are given n cities in a fixed sequence. Each city has a treasure value t[i] and a colour c[i]. You move left to right, forward only, and decide whether to visit or skip each city. When you visit a city, you collect:

• a * t[i] if the colour of the previously visited city equals c[i] (same colour)
• b * t[i] if the previously visited city had a different colour, or this is the first city you visit

The values a, b, and all entries in t[] can be negative. Colours in c[] range from 1 to n. You may skip all cities (profit = 0). Find the maximum profit.

Canonical example (from placement tests):

• a = 5, b = -7
• t = {4, 8, 2, 9}, c = {2, 2, 3, 2}

Optimal path: cities 1 → 2 → 4, profit = 57. The derivation is in the worked-example section below.

This problem type appears in coding rounds at analytics and quant-oriented firms. The key constraint that makes it non-trivial: a and b can be negative, meaning visiting a city in the wrong context subtracts from your total. That rules out any strategy that visits every city unconditionally.

Why Greedy Selection Fails

A naive approach might say: always visit the city with the highest treasure, or always prefer same-colour cities. Both strategies fail on the canonical example.

Consider visiting only city 4 (highest treasure = 9):

• City 4 is the first visited, so profit = b * 9 = -7 * 9 = -63.

Consider visiting all four cities:

• City 1 (first): b * 4 = -28
• City 2 (same colour as city 1): a * 8 = 40
• City 3 (different colour from city 2): b * 2 = -14
• City 4 (different colour from city 3): b * 9 = -63
• Total = -65

The correct answer of 57 comes from skipping city 3, which breaks the same-colour chain that connects city 2 and city 4. Greedy cannot see that skipping an intermediate city improves the rewards for a later one.

This is the standard failure mode that dynamic programming is designed to address: the locally optimal action is not globally optimal.

The DP Approach

Define dp[i] as the maximum profit achievable when city i is the last visited city.

Base Case

City i is visited as the first city (no predecessor):

• dp[i] = b * t[i]

Transition

For each city j where j < i:

• If c[j] == c[i]: dp[i] = max(dp[i], dp[j] + a * t[i])
• If c[j] != c[i]: dp[i] = max(dp[i], dp[j] + b * t[i])

Final Answer

answer = max(0, dp[1], dp[2], ..., dp[n])

The floor of 0 covers the case where every possible city visit produces negative profit.

Edge Cases

Two edge cases worth checking before running the main loop:

• Single city (n=1): dp[0] = b * t[0]; answer = max(0, b * t[0]).
• All cities the same colour: every transition uses the a multiplier. If a is positive, visit all cities; if a is negative, visit none (answer = 0).

The structure is analogous to the 0/1 knapsack problem: for each item (city), you decide to include or exclude it, and the reward depends on context (colour of the previous inclusion rather than weight capacity).

Worked Example, Step by Step

Input: a = 5, b = -7, t = {4, 8, 2, 9}, c = {2, 2, 3, 2}

Work through the DP table row by row:

• dp[1]: base case only. dp[1] = b * t[1] = -7 * 4 = -28

• dp[2]: base case dp[2] = -7 * 8 = -56. Then check j=1: c[1]=2 equals c[2]=2, same colour. dp[2] = max(-56, dp[1] + a * t[2]) = max(-56, -28 + 40) = max(-56, 12) = 12

• dp[3]: base case dp[3] = -7 * 2 = -14.

  • j=1: c[1]=2 vs c[3]=3, different. dp[3] = max(-14, dp[1] + b * t[3]) = max(-14, -28 + (-14)) = max(-14, -42) = -14
  • j=2: c[2]=2 vs c[3]=3, different. dp[3] = max(-14, dp[2] + b * t[3]) = max(-14, 12 + (-14)) = max(-14, -2) = -2

• dp[4]: base case dp[4] = -7 * 9 = -63.

  • j=1: c[1]=2 equals c[4]=2, same. dp[4] = max(-63, dp[1] + a * t[4]) = max(-63, -28 + 45) = max(-63, 17) = 17
  • j=2: c[2]=2 equals c[4]=2, same. dp[4] = max(17, dp[2] + a * t[4]) = max(17, 12 + 45) = max(17, 57) = 57
  • j=3: c[3]=3 vs c[4]=2, different. dp[4] = max(57, dp[3] + b * t[4]) = max(57, -2 + (-63)) = max(57, -65) = 57

Final dp array: {-28, 12, -2, 57}

Answer = max(0, -28, 12, -2, 57) = 57

The optimal path is cities 1 → 2 → 4:

• City 1 (first visit): b * 4 = -28
• City 2 (same colour 2 as city 1): a * 8 = 40
• City 4 (same colour 2 as city 2): a * 9 = 45
• Total = -28 + 40 + 45 = 57 ✓

Python Implementation

def treasure_cities(a, b, t, c):
    """
    Returns the maximum profit from the Treasure and Cities problem.

    Parameters:
        a (int): multiplier for same-colour visit
        b (int): multiplier for different-colour visit (or first city)
        t (list): treasure values for each city
        c (list): colour of each city

    Returns:
        int: maximum profit (0 if all visits produce negative profit)
    """
    n = len(t)
    dp = [0] * n

    # Base case: each city visited as first (no predecessor)
    for i in range(n):
        dp[i] = b * t[i]

    # Transition: try every predecessor j for each city i
    for i in range(1, n):
        for j in range(i):
            if c[j] == c[i]:
                dp[i] = max(dp[i], dp[j] + a * t[i])
            else:
                dp[i] = max(dp[i], dp[j] + b * t[i])

    return max(0, max(dp))


# Canonical example
a, b = 5, -7
t = [4, 8, 2, 9]
c = [2, 2, 3, 2]
print(treasure_cities(a, b, t, c))  # Output: 57

Complexity and Variants

• Time complexity: O(n²) — the outer loop runs n times; the inner loop runs up to n-1 times.
• Space complexity: O(n) — only the dp array is stored.

For placement coding rounds, O(n²) with n in the low hundreds is fast enough. The inner loop can be optimised by separating cities into colour buckets and tracking the running maximum per colour, but the two-loop version is what most interviewers expect at the screening stage.

Related DP patterns that share the same “visit or skip with context-dependent reward” structure:

• 0/1 knapsack (capacity constraint replaces colour constraint)
• Longest Increasing Subsequence (order constraint, no colour)
• Maximum sum of non-adjacent elements (adjacency constraint)

DP problems of this type appear in the coding section of placement tests at analytics and product firms. The D.E. Shaw placement process, for instance, includes a dynamic programming coding problem in its online round. Working through problems like Treasure and Cities builds the pattern-recognition needed for those rounds.

For broader placement-test preparation, the campus placement evaluation test guide covers how aptitude and coding sections are structured across companies. If you’re building your DP foundation from books, the best books for placement preparation article lists the standard resources by section.

The same chain-optimisation logic that makes this problem interesting is the basis of how sequence models make decisions: a model choosing its next token weighs prior context the same way this algorithm weighs prior colour. TinkerLLM at ₹299 is an LLM playground where you can run those experiments without a GPU setup. The DP intuition you’ve built tracing through dp[4] = 57 carries directly into understanding why context length changes model behaviour.