ArrayOrInt.java

## ArrayOrInt.java
package flatten;

import java.util.Arrays;
import java.util.Optional;
import java.util.OptionalInt;
import java.util.function.Function;
import java.util.function.IntFunction;
import java.util.stream.IntStream;

public class ArrayOrInt {
	private final Optional<ArrayOrInt[]> array;
	private final OptionalInt integer;

	public ArrayOrInt(Optional<ArrayOrInt[]> array, OptionalInt integer) {
		if (!(array.isPresent() ^ integer.isPresent())) {
			throw new IllegalArgumentException();
		}
		this.array = array;
		this.integer = integer;
	}

	public static ArrayOrInt of(int value) {
		return new ArrayOrInt(Optional.empty(), OptionalInt.of(value));
	}

	public static ArrayOrInt of(int... values) {
		ArrayOrInt[] objects = IntStream.of(values).mapToObj(ArrayOrInt::of).toArray(ArrayOrInt[]::new);
		return ArrayOrInt.of(objects);
	}

	public static ArrayOrInt of(ArrayOrInt... values) {
		return new ArrayOrInt(Optional.of(values), OptionalInt.empty());
	}

	public <T> T map(Function<ArrayOrInt[], T> ifArray, IntFunction<T> ifInt) {
		return array.map(ifArray).orElseGet(() -> ifInt.apply(integer.getAsInt()));
	}

	@Override
	public String toString() {
		return map(Arrays::toString, Integer::toString);
	}
}

## Flatten.java
package flatten;

import static flatten.ArrayOrInt.of;

import java.util.Arrays;
import java.util.stream.IntStream;
import java.util.stream.Stream;

public class Flatten {

	public static void main(String[] args) {
		ArrayOrInt input = of(of(1), of(2, 3), of(of(4), of(5, 6)));

		System.out.println(input);
		// [1, [2, 3], [4, [5, 6]]]

		int[] flattened = flatten(input).toArray();
		System.out.println(Arrays.toString(flattened));
		// [1, 2, 3, 4, 5, 6]
	}

	private static IntStream flatten(ArrayOrInt in) {
		return in.map(items -> Stream.of(items).flatMapToInt(Flatten::flatten), IntStream::of);
	}
}