CtrlAltDileep

How to Solve Two Sum in Java

Written by

Dileep

in

, ,

Today we’ll be tackling the quintessential problem everyone starts with on LeetCode: Two Sum. It’s approachable but teaches a genuinely useful lesson: how to  pick the right data structure for a problem.

The Problem

The problem from LeetCode looks like this:

Given an array of integer nums and an integer target, return the indices of the two numbers such that they add up to target.

You may assume that each input would have exactly one solution, and you may not use the same element twice. You can return the answer in any order.

Example 1:

Input: nums [3, 2, 4] and target = 6

Output: [1,2]

Explanation: because nums[1] + nums[2] == 6 so we return [1, 2]

Example 2:

Input: nums [3,3] and target = 6

Output: [0,1]

Explanation: because nums[0] + nums[1] == 6 so we return [0,1]

From the examples provided we can learn two new things about the problem. Firstly the numbers in the input are unsorted and that although we cannot use the same number twice a number can appear multiple times in the input.

Brute Force

I like to try to solve problems using a naive brute-force approach first. It often gives insight into how to give the optimal solution for the problem later. So what does a brute force algorithm for this look like? Well we can try computing the sum of every pair of numbers and seeing if they equal our target number. We can pseudocode an algorithm like below.

  1. For each number in the input
    1. Go through every other number in the input
      1. Add the current number and the other number up
      2. If their sum equals target then return their indexes
      3. Otherwise go to the next pair of numbers
public int[] twoSum(int[] nums, int target) {
   // go through each number in the input nums
   for(int i = 0; i < nums.length; i++) {
       // go through each other number in the input
       for(int j = i + 1; j < nums.length; j++) {
           // if the sum of each number is equal to target then return the indexes
           if(nums[i] + nums[j] == target) {
               return new int[]{i, j};
           }
           // Otherwise go to the next pair of numbers
       }
   }

   // we should never reach this but if we do then we got a bad input
   throw new IllegalArgumentException();

}

With that coded up let’s take a look at the time and space complexity of our work. In terms of space we don’t allocate any new memory so the solution is  O(1) in space. When it comes to time complexity we are O(n2). The rationale for the time complexity is that the worst case scenario for this algorithm is that the solution is the last two numbers in the input array. In that situation we would have to loop through the entire array for each preceding number until we reach the last two.

Key Insight for Optimization

The key to coming up with an optimal solution actually comes from looking at the brute-force approach. If we analyze it we notice that we lose a lot of time checking each individual pair in the inner for loop. If we could do that more efficiently then the algorithm would become dramatically faster. 

Turns out there is actually a data structure that could help with this! We could introduce a hashmap to look up numbers and indices efficiently. Using that we can avoid the extra for-loop and check immediately if we have a number that adds up to target or not. How do we know what key to look up? Well for any number and a target number we know that we need to find something equal to target - current number to have a solution. So we can look up that target - current value in our HashMap quickly.

Optimal Solution

So what does the optimal solution look like? The algorithm would look something like this. 

  1. We create a map that stores values to their indexes
  2. We go through each number in the input array
    1. We calculate the number we need to add our current number to get the target number
    2. We check if our map has that number we need
      1. If yes then return the indexes and we are done!
      2. Otherwise add the current number and index to our map and move on

So lets try coding that out.

public int[] twoSum(int[] nums, int target) {
   // We create a map that stores values to their indexes
   Map<Integer, Integer> valToIndex = new HashMap<>();


   // We go through each number in the input array
   for (int i = 0; i < nums.length; i++) {
       // We calculate the number we need to add our current number to get the target number
       int goal = target - nums[i];
        //We check if our map has that number we need
       if (valToIndex.containsKey(goal)) {
           // If yes then return the indexes and we are done!
           return new int[]{i, valToIndex.get(goal)};
       }


       // Otherwise add the current number and index to our map and move on
       valToIndex.put(nums[i], i);
   }


   // we should never reach this but if we do then we got a bad input
   throw new IllegalArgumentException();
}

With this new approach we have actually improved our time complexity to O(n). This is because the look ups in the hashmap are O(1) and we only have to loop through all the numbers in the input array once in the worst case scenario that the two numbers we need are at the end of the input array. However this speed improvement did not come for free. Out space complexity actually got worse and became O(n). This is because we need to allocate a Map that could potentially store our entire input array!

Key Takeaway

I think the most valuable lesson the Two Sum problem teaches is that if you are repeatedly searching the same list of things over and over we can make that more efficient by using a HashMap lookup pattern. This technique can be used in other LeetCode problems like Four Sum and Three Sum. So hopefully this article prepares you for the next challenges ahead!

Also if you want more of a walk through check out this Youtube Video of me solving the problem.

More posts