Method cut

models/treebased/data/data_cutter.py:28–56

(self)

        self._number = number

    def cut(self):
        user_behav = dict()
        user_ids = list()
        with open(self._input) as f:
            for line in f:
                arr = line.strip().split(',')
                if len(arr) != 5:
                    break

                if arr[0] not in user_behav:
                    user_ids.append(arr[0])
                    user_behav[arr[0]] = list()

                user_behav[arr[0]].append(line)

        random.shuffle(user_ids)
        test_user_ids = user_ids[:self._number]
        train_user_ids = user_ids[self._number:]

        # write train data set
        with open(self._train, 'w') as f:
            for uid in train_user_ids:
                for line in user_behav[uid]:
                    f.write(line)

        with open(self._test, 'w') as f:
            for uid in test_user_ids:
                for line in user_behav[uid]:
                    f.write(line)


if __name__ == '__main__':