Number of comments per post

# Table: Submissions
#
# +---------------+----------+
# | Column Name   | Type     |
# +---------------+----------+
# | sub_id        | int      |
# | parent_id     | int      |
# +---------------+----------+
# This table may have duplicate rows.
# Each row can be a post or comment on the post.
# parent_id is null for posts.
# parent_id for comments is sub_id for another post in the table.
#
#
#
# Write a solution to find the number of comments per post. The result table should contain post_id and its corresponding number_of_comments.
#
# The Submissions table may contain duplicate comments. You should count the number of unique comments per post.
#
# The Submissions table may contain duplicate posts. You should treat them as one post.
#
# The result table should be ordered by post_id in ascending order.
#
# The result format is in the following example.
#
# Example 1:
# Input:
# Submissions table:
# +---------+------------+
# | sub_id  | parent_id  |
# +---------+------------+
# | 1       | Null       |
# | 2       | Null       |
# | 1       | Null       |
# | 12      | Null       |
# | 3       | 1          |
# | 5       | 2          |
# | 3       | 1          |
# | 4       | 1          |
# | 9       | 1          |
# | 10      | 2          |
# | 6       | 7          |
# +---------+------------+
# Output:
# +---------+--------------------+
# | post_id | number_of_comments |
# +---------+--------------------+
# | 1       | 3                  |
# | 2       | 2                  |
# | 12      | 0                  |
# +---------+--------------------+
# Explanation:
# The post with id 1 has three comments in the table with id 3, 4, and 9. The comment with id 3 is repeated in the table, we counted it only once.
# The post with id 2 has two comments in the table with id 5 and 10.
# The post with id 12 has no comments in the table.
# The comment with id 6 is a comment on a deleted post with id 7 so we ignored it.

import pandas as pd


def count_comments(submissions: pd.DataFrame) -> pd.DataFrame:
  # Posts are rows where parent_id is null
  posts = (
    submissions[submissions["parent_id"].isna()][["sub_id"]]
    .drop_duplicates()
    .rename(columns={"sub_id": "post_id"})
  )

  # Comments are rows where parent_id is not null
  comments = submissions[submissions["parent_id"].notna()][
    ["sub_id", "parent_id"]
  ].drop_duplicates()

  # Count unique comments per post
  comment_counts = comments.groupby("parent_id")["sub_id"].count().reset_index()
  comment_counts.columns = ["post_id", "number_of_comments"]

  result = posts.merge(comment_counts, on="post_id", how="left").fillna(0)
  result["number_of_comments"] = result["number_of_comments"].astype(int)
  return result.sort_values("post_id").reset_index(drop=True)