Solution to Leetcode's Contiguous Array

This is a medium Leetcode problem 525. See below image for description:

ALGORITHM

Create a map that stores the current count to the index.
Loop through the array, if the current value is 0, subtract 1 from count. if it is 1, add 1 to count.
if the map does not have the current value of count as a key, put the new value in map with key as the current count and the value as the index.
if the map has the current count as key, get the maximum between the current max and (index - value of the key).
return max.

CODE

 public int findMaxLength(int[] nums) {
        int count = 0;
        int max = 0;
        HashMap<Integer, Integer> map = new HashMap<>();
        map.put(0, -1);

        for (int i = 0; i < nums.length; i++) {
            if (nums[i] == 0)
                count += -1;
            else
                count += 1;

            if (map.containsKey(count))
                max = Math.max(max, i - map.get(count));
            else
                map.put(count, i);

        }

        return max;
    }