javaarraysmatrixmultidimensional-arrayjava-stream

Stream through a portion of a 2D array in Java


Let's say you have a matrix [[1,2,5],[3,4,6]] and you need to work with the first two elements of each row together [1,2,3,4]. You don't want to store this in memory or use nested loops. Is there a way to achieve this using streams in Java?


Solution

  • Here is example

    int[][] arr = {{1, 2, 5}, {3, 4, 6}};
    
    Arrays.stream(arr)
            .flatMapToInt(a -> Arrays.stream(a).limit(2))
            .forEach(System.out::println);
    

    EDIT Here is summary of all the methods:

    @Fork(value = 1, warmups = 1)
    @BenchmarkMode(Mode.AverageTime)
    @State(Scope.Benchmark)
    public class MainTest {
        public static final int FIRST_SIZE = 1000;
        private static final int SECOND_SIZE = 100000;
        private int[][] arr;
    
        @Setup(Level.Trial)
        public void setUp() {
            arr = IntStream.range(0, FIRST_SIZE)
                    .mapToObj(x -> new Random().ints(SECOND_SIZE).toArray())
                    .toArray(int[][]::new);
        }
    
        @Benchmark
        public void loop(Blackhole blackhole) {
            int sum = 0;
            for (int i = 0; i < arr.length; i++) {
                for (int j = 0; j < arr[i].length; j++) {
                    sum += arr[i][j];
                }
            }
            blackhole.consume(sum);
        }
    
        @Benchmark
        public void singleLoop(Blackhole blackhole) {
            int sum = 0;
            for (int i = 0; i < FIRST_SIZE * SECOND_SIZE; i++) {
                sum += arr[i / SECOND_SIZE][i % SECOND_SIZE];
            }
            blackhole.consume(sum);
        }
    
        @Benchmark
        public void enhancedLoop(Blackhole blackhole) {
            int sum = 0;
            for (int[] ints : arr) {
                for (int anInt : ints) {
                    sum += anInt;
                }
            }
            blackhole.consume(sum);
        }
    
        @Benchmark
        public void stream(Blackhole blackhole) {
            int sum = Arrays.stream(arr)
                    .flatMapToInt(a -> Arrays.stream(a))
                    .sum();
            blackhole.consume(sum);
        }
    
        @Benchmark
        public void parallelStream(Blackhole blackhole) {
            int sum = Arrays.stream(arr)
                    .parallel()
                    .flatMapToInt(a -> Arrays.stream(a))
                    .sum();
            blackhole.consume(sum);
        }
    
    
        @Benchmark
        public void mapMultiToInt(Blackhole blackhole) {
            int sum = Arrays.stream(arr).mapMultiToInt((row, consumer) -> {
                for (int i = 0; i < row.length; i++) {
                    consumer.accept(row[i]);
                }
            }).sum();
    
            blackhole.consume(sum);
        }
    
        public static void main(String[] args) throws IOException {
            org.openjdk.jmh.Main.main(args);
        }
    }
    

    And here is result on my machine:

    Benchmark                       Mode  Cnt  Score   Error  Units
    MainTest.enhancedLoop           avgt   15  0.047 ± 0.003   s/op
    MainTest.loop                   avgt   15  0.063 ± 0.012   s/op
    MainTest.mapMultiToInt          avgt   15  0.060 ± 0.006   s/op
    MainTest.parallelStream         avgt   15  0.047 ± 0.003   s/op
    MainTest.singleLoop             avgt   15  0.297 ± 0.020   s/op
    MainTest.stream                 avgt   15  0.051 ± 0.005   s/op