Line data    Source code

       1              : // Copyright (c) 2024 The Authors. All rights reserved.
       2              : //
       3              : // Licensed under the Apache License, Version 2.0 (the "License");
       4              : // you may not use this file except in compliance with the License.
       5              : // You may obtain a copy of the License at
       6              : //
       7              : //      https://www.apache.org/licenses/LICENSE-2.0
       8              : //
       9              : // Unless required by applicable law or agreed to in writing, software
      10              : // distributed under the License is distributed on an "AS IS" BASIS,
      11              : // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
      12              : // See the License for the specific language governing permissions and
      13              : // limitations under the License.
      14              : 
      15              : // Authors: liubang (it.liubang@gmail.com)
      16              : 
      17              : #include <gtest/gtest.h>
      18              : 
      19              : #include <iostream>
      20              : #include <vector>
      21              : 
      22              : namespace {
      23              : class NestedInteger {
      24              : public:
      25            8 :     NestedInteger(int value) : value_(value), is_integer_(true) {}
      26              : 
      27            1 :     NestedInteger(const std::vector<NestedInteger>& list) : is_integer_(false), list_(list) {}
      28              : 
      29            3 :     NestedInteger(const std::vector<int>& ints) {
      30            8 :         for (int i : ints) {
      31            5 :             list_.emplace_back(i);
      32              :         }
      33            3 :         is_integer_ = false;
      34            3 :     }
      35              : 
      36            3 :     NestedInteger(const NestedInteger&) = default;
      37              :     NestedInteger& operator=(const NestedInteger&) = default;
      38            7 :     NestedInteger(NestedInteger&&) = default;
      39              : 
      40           12 :     bool isInteger() const { return is_integer_; }
      41              : 
      42            8 :     int getInteger() const { return value_; }
      43              : 
      44            4 :     const std::vector<NestedInteger>& getList() const { return list_; }
      45              : 
      46              : private:
      47              :     int value_{0};
      48              :     bool is_integer_{true};
      49              :     std::vector<NestedInteger> list_;
      50              : };
      51              : 
      52              : class NestedIterator {
      53              : public:
      54            2 :     NestedIterator(const std::vector<NestedInteger>& nestedList) { visit(nestedList); }
      55              : 
      56            8 :     int next() { return list_[idx_++]; }
      57              : 
      58           10 :     bool hasNext() { return idx_ < list_.size(); }
      59              : 
      60              : private:
      61            6 :     void visit(const std::vector<NestedInteger>& nestedList) {
      62           18 :         for (auto& nestedInteger : nestedList) {
      63           12 :             if (nestedInteger.isInteger()) {
      64            8 :                 list_.push_back(nestedInteger.getInteger());
      65              :             } else {
      66            4 :                 visit(nestedInteger.getList());
      67              :             }
      68              :         }
      69            6 :     }
      70              : 
      71              : private:
      72              :     std::vector<int> list_;
      73              :     int idx_{0};
      74              : };
      75              : } // namespace
      76              : 
      77            4 : TEST(Leetcode, flatten_nested_list_iterator) {
      78              :     {
      79              :         // [[1,1],2,[1,1]]
      80            1 :         std::vector<NestedInteger> inputs;
      81            1 :         std::vector<int> output;
      82            1 :         std::vector<int> exp = {1, 1, 2, 1, 1};
      83            1 :         inputs.emplace_back(std::vector<int>{1, 1});
      84            1 :         inputs.emplace_back(2);
      85            1 :         inputs.emplace_back(std::vector<int>{1, 1});
      86              : 
      87            1 :         NestedIterator it(inputs);
      88            6 :         while (it.hasNext()) {
      89            5 :             output.push_back(it.next());
      90              :         }
      91            1 :         EXPECT_EQ(exp, output);
      92            1 :     }
      93              : 
      94              :     {
      95              :         // [1,[4,[6]]]
      96            1 :         std::vector<NestedInteger> inputs;
      97            1 :         std::vector<int> output;
      98            1 :         std::vector<int> exp = {1, 4, 6};
      99            1 :         inputs.emplace_back(1);
     100              : 
     101            1 :         std::vector<NestedInteger> inner;
     102            1 :         inner.emplace_back(4);
     103            1 :         inner.emplace_back(std::vector<int>{6});
     104            1 :         inputs.emplace_back(inner);
     105              : 
     106            1 :         NestedIterator it(inputs);
     107            4 :         while (it.hasNext()) {
     108            3 :             output.push_back(it.next());
     109              :         }
     110            1 :         EXPECT_EQ(exp, output);
     111            1 :     }
     112            1 : }