小红书2018提前批笔试题之0、1个数相同的最长子序列

package redbook;

/**
 * 给定一个只包含0和1的数组，找到其中包含相同0的个数和1的个数的最长子序，输出子序列的长度，要求在o（n）时间内完成。
 */

import java.util.HashMap;
import java.util.Map;
import java.util.Scanner;
import java.util.TreeSet;

public class Main2 {
    public static void main(String[] args) {
        Scanner in = new Scanner(System.in);
        while (in.hasNext()) {
            int n = in.nextInt();
            //定义数组，存储输入的01串
            int[] arr = new int[n];
            //输入01数字串，并且把其中的0换成-1
            for (int i = 0; i < n; i++) {
                arr[i] = in.nextInt() == 0 ? -1 : 1;
            }
            //调用方法,输出包含相同0的个数和1的个数的最长子序的长度
            int maxLength = maxLength(arr, arr.length);
            System.out.println(maxLength);
        }
        in.close();
    }

    public static int maxLength(int[] arr, int arrLength) {
        int max=0;//保存最大的长度
        int sum = 0;//保存累加的和
        int index = 0;
        Map<Integer, TreeSet<Integer>> hashMap = new HashMap<>();//保存数字的和与其位置的映射

        for (int i = 0; i < arrLength; i++) {
            sum = sum + arr[i];//依次累加求和

            //第一种情况：子串都是从头开始
            if (sum == 0) {
                index = i;
            }

            //第二种情况：子串都不是从头开始，将数字的和与其位置的映射存入HashMap中
            if (hashMap.containsKey(sum)) {
                hashMap.get(sum).add(i);
            } else {
                TreeSet<Integer> treeSet = new TreeSet<>();
                treeSet.add(i);
                hashMap.put(sum, treeSet);
            }
        }

        //判断
        if (index != 0) {
            max = index + 1;//第一种情况：子串都是从头开始,子串最大值为index+1
        } else { //第二种情况：子串不是从头开始,子串最大值为位置的最大值-位置的最小值
            for (TreeSet<Integer> treeSet : hashMap.values()) {
                max = (treeSet.last() - treeSet.first()) > max ? (treeSet.last() - treeSet.first()) : max;
            }
        }

        return max;
    }

}