I was solving this problem and after I solved it by myself I decided to implement editorial's solution. Strangely my first submission was really slow 55982450, but after changing the order of dimensions it got 5x faster 55999256.